How to load PluggableDevice into tensorflow to use Mac M1 GPU? - tensorflow

Question

I have successfully built tensorflow v2.9 C library by following this gist on Mac M1.
When I check the devices in Rust code, only CPU device is available. bundle.session.device_list() only returns CPU device:
let bundle = SavedModelBundle::load(
&SessionOptions::new(),
&["serve"],
&mut graph,
export_dir
).expect("Unable to load model from disk");
println!("{:?}", bundle.session.device_list() )
Device { name: "/job:localhost/replica:0/task:0/device:CPU:0", device_type: "CPU", memory_bytes: 268435456, incarnation: 10072007419359857694 }]
The Rust code are using bindings to tensorflow's C API (e.g. device_list = TF_DeviceList). Apple M1's GPU is supported by MPS plugin (https://developer.apple.com/metal/tensorflow-plugin/). From my test, it does work in python. The MPS is implemented as a PluggableDevice to tensorflow, that is to say, it can be loaded without modification to tensorflow source code.
How can I use the deivce plugin in tensorflow?

Answer 1

There is an experimental API called TF_LoadPluggableDeviceLibrary
So first ensure Apple's metal plugin is installed, then find the location of file libmetal_plugin.dylib. Either copy this file into PATH, or add its directory into DYLD_LIBRARY_PATH environment variable.
On application startup, before loading the model, load this library first:
pub fn load_plugable_device(library_filename: &str) -> Result<(), Status> {
use std::ffi::{ CString };
let c_filename = CString::new(library_filename)?;
let raw_lib = unsafe {
let raw_status: *mut tf::TF_Status = tf::TF_NewStatus();
let raw_lib = tf::TF_LoadPluggableDeviceLibrary(c_filename.as_ptr(), raw_status);
if !raw_status.is_null() {
tf::TF_DeleteStatus(raw_status)
}
raw_lib
};
if raw_lib.is_null() {
Err(Status::new())
} else {
Ok(())
}
}
I am using Rust bindings and this experimental API is missing in tensorflow-sys crate. Hence I add it there. You may check my github repository for source code.
Then the lovely GPU is enabled.
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.