use std::{cmp, collections::HashMap, convert::TryFrom, iter::FromIterator, mem, str};
use nom::{alpha1, digit1, le_i32, le_i64, le_u16, le_u32, le_u64, le_u8, types::CompleteStr};
use serde;
use serde_json;
use super::{DLTensor, DataType, Module, Storage, TVMContext, Tensor};
use crate::{
common::value::TVMArgValue,
errors::{Error, ErrorKind, Result},
ffi::runtime::{DLDataTypeCode_kDLFloat, DLDataTypeCode_kDLInt, DLDataTypeCode_kDLUInt},
};
// @see `kTVMNDArrayMagic` in `ndarray.h`
const _NDARRAY_MAGIC: u64 = 0xDD5E40F096B4A13F;
// @see `kTVMNDArrayListMagic` in `graph_runtime.h`
const _NDARRAY_LIST_MAGIC: u64 = 0xF7E58D4F05049CB7;
/// A TVM computation graph.
///
/// # Examples
///
/// ```
/// let graph_json = fs::read_to_string("graph.json")).unwrap();
/// let graph = Graph::try_from(&graph_json).unwrap();
/// ```
#[derive(Serialize, Deserialize, Debug)]
pub struct Graph {
pub nodes: Vec<Node>,
pub arg_nodes: Vec<usize>,
pub heads: Vec<Entry>,
pub node_row_ptr: Option<Vec<usize>>,
pub attrs: Option<HashMap<String, serde_json::Value>>,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct Entry {
pub id: usize,
pub index: usize,
pub version: usize,
}
impl Graph {
fn entry_index(&self, entry: &Entry) -> Result<usize> {
self.node_row_ptr
.as_ref()
.map(|nrp| nrp[entry.id] + entry.index)
.ok_or("Missing node_row_ptr.".into())
}
/// Attempt to deserialize a JSON attribute to a type `T`.
fn get_attr<T: serde::de::DeserializeOwned>(&self, attr: &str) -> Result<T> {
Ok(serde_json::from_value::<T>(
self.attrs
.as_ref()
.ok_or(ErrorKind::GraphFormatError(
"Missing graph attrs".to_string(),
))?
.get(attr)
.ok_or(ErrorKind::GraphFormatError(format!(
"Missing {} attr",
attr
)))?
.to_owned(),
)?)
}
}
#[derive(Serialize, Deserialize, Debug)]
pub struct Node {
pub op: String,
pub name: String,
pub inputs: Vec<Entry>,
pub attrs: Option<HashMap<String, String>>,
pub control_deps: Option<Vec<Entry>>,
}
struct NodeAttrs {
func_name: String,
num_outputs: usize,
flatten_data: bool,
}
impl Node {
fn parse_attrs(&self) -> Result<NodeAttrs> {
let attrs = self
.attrs
.as_ref()
.ok_or(format!("Missing node.attrs for `{}`", self.name))?;
let func_name = attrs
.get("func_name")
.ok_or(format!("Node `{}` is missing attrs.func_name", self.name))?
.to_string();
let num_outputs = attrs
.get("num_outputs")
.ok_or(format!("Node `{}` is missing attrs.num_outputs", self.name))?
.parse::<usize>()?;
let flatten_data = attrs
.get("flatten_data")
.ok_or(format!(
"Node `{}` is missing attrs.flatten_data",
self.name
))?
.parse::<u8>()?
== 1;
Ok(NodeAttrs {
func_name,
num_outputs,
flatten_data,
})
}
}
impl<'a> TryFrom<&'a String> for Graph {
type Error = Error;
fn try_from(graph_json: &String) -> Result<Self> {
let graph = serde_json::from_str(graph_json)?;
Ok(graph)
}
}
impl<'a> TryFrom<&'a str> for Graph {
type Error = Error;
fn try_from(graph_json: &'a str) -> Result<Self> {
let graph = serde_json::from_str(graph_json)?;
Ok(graph)
}
}
/// A executor for a TVM computation graph.
///
/// # Examples
///
/// ```
/// use ndarray::Array;
///
/// let syslib = SystemLibModule::default(); // a provider of TVM functions
///
/// let mut params_bytes = Vec::new();
/// fs::File::open("graph.params").unwrap().read_to_end(&mut params_bytes).unwrap();
/// let params = tvm::runtime::load_param_dict(¶ms_bytes).unwrap();
///
/// let graph = Graph::try_from(&fs::read_to_string("graph.json").unwrap()).unwrap();
///
/// let mut exec = GraphExecutor::new(graph, &syslib).unwrap();
/// exec.load_params(params);
///
/// let x = Array::from_vec(vec![1f32, 2., 3., 4.]);
/// exec.set_input("data", x.into());
/// exec.run();
/// let output = exec.get_output(0).unwrap();
///
/// println!("{:#?}", Array::try_from(output).unwrap());
/// ```
pub struct GraphExecutor<'m, 't> {
graph: Graph,
op_execs: Vec<Box<Fn() + 'm>>,
tensors: Vec<Tensor<'t>>,
}
unsafe impl<'m, 't> Send for GraphExecutor<'m, 't> {}
impl<'m, 't> GraphExecutor<'m, 't> {
pub fn new<M: 'm + Module>(graph: Graph, lib: &'m M) -> Result<Self> {
let tensors = Self::setup_storages(&graph)?;
Ok(GraphExecutor {
op_execs: Self::setup_op_execs(&graph, lib, &tensors)?,
tensors: tensors,
graph: graph,
})
}
/// Runs the computation graph.
pub fn run(&self) {
self.op_execs.iter().for_each(|op_exec| {
op_exec();
});
}
/// Allocates `Storages` for each `storage_id` and returns `Tensor`s to hold each output.
fn setup_storages<'a>(graph: &'a Graph) -> Result<Vec<Tensor<'t>>> {
let storage_ids = graph.get_attr::<(String, Vec<usize>)>("storage_id")?.1;
let shapes = graph.get_attr::<(String, Vec<Vec<i64>>)>("shape")?.1;
let dtypes = graph
.get_attr::<(String, Vec<String>)>("dltype")?
.1
.iter()
.map(|dltype| {
if let Ok((_, dtype)) = tvm_str_to_type(CompleteStr(dltype)) {
Ok(dtype)
} else {
Err(ErrorKind::GraphFormatError(
format!("Invalid dltype: {}", dltype).to_string(),
)
.into())
}
})
.collect::<Result<Vec<DataType>>>()?;
let align = dtypes.iter().map(|dtype| dtype.bits as usize).max();
let mut storage_num_bytes = vec![0usize; *storage_ids.iter().max().unwrap_or(&1) + 1];
for (i, &storage_id) in storage_ids.iter().enumerate() {
let dtype_size = dtypes[i].bits * dtypes[i].lanes >> 3;
let nbytes = dtype_size * shapes[i].iter().product::<i64>() as usize;
storage_num_bytes[storage_id] = cmp::max(nbytes, storage_num_bytes[storage_id]);
}
let mut storages: Vec<Storage> = storage_num_bytes
.into_iter()
.map(|nbytes| Storage::new(nbytes, align))
.collect::<Result<Vec<Storage>>>()?;
let tensors = izip!(storage_ids, shapes, dtypes)
.map(|(storage_id, shape, dtype)| {
let storage = storages[storage_id].view();
Tensor {
data: mem::replace(&mut storages[storage_id], storage),
ctx: TVMContext::default(),
dtype: dtype,
size: shape.iter().product::<i64>() as usize,
shape: shape,
strides: None,
byte_offset: 0,
}
})
.collect();
Ok(tensors)
}
/// Creates closures which represent the computation performed by this graph.
fn setup_op_execs<M: 'm + Module>(
graph: &Graph,
lib: &'m M,
tensors: &Vec<Tensor<'t>>,
) -> Result<Vec<Box<Fn() + 'm>>> {
ensure!(graph.node_row_ptr.is_some(), "Missing node_row_ptr.");
let node_row_ptr = graph.node_row_ptr.as_ref().unwrap();
let mut op_execs = Vec::new();
for (i, node) in graph.nodes.iter().enumerate() {
if node.op == "null" {
continue;
}
ensure!(node.op == "tvm_op", "Only TVM ops are supported.");
ensure!(node.attrs.is_some(), "Missing node attrs.");
let attrs = node.parse_attrs()?;
if attrs.func_name == "__nop" {
continue;
}
let func = lib
.get_function(&attrs.func_name)
.ok_or(format!("Missing function {}", attrs.func_name))?;
let arg_indices = node
.inputs
.iter()
.map(|entry| graph.entry_index(entry))
.chain((0..attrs.num_outputs).map(|oi| Ok(node_row_ptr[i].clone() + oi)));
let dl_tensors = arg_indices
.map(|idx| {
let tensor = &tensors[idx?];
Ok(if attrs.flatten_data {
DLTensor::from_tensor(tensor, true /* flatten */)
} else {
DLTensor::from(tensor)
})
})
.collect::<Result<Vec<DLTensor>>>()
.unwrap();
let op: Box<Fn()> = box move || {
let args = dl_tensors
.iter()
.map(|t| t.into())
.collect::<Vec<TVMArgValue>>();
func(args.as_slice());
};
op_execs.push(op);
}
Ok(op_execs)
}
pub fn load_params(&mut self, params: HashMap<String, Tensor>) {
params.into_iter().for_each(|(name, param)| {
self.set_input(name, param);
})
}
pub fn set_input<S: AsRef<str>>(&mut self, name: S, value: Tensor) {
if let Some(idx) = self.get_input_index(name.as_ref()) {
// TODO: consider `new_with_params` to avoid ever allocating
let ptr = self.tensors[idx].data.as_ptr();
let mut to_replace = self.tensors.iter_mut().filter(|t| t.data.as_ptr() == ptr);
let mut owner = to_replace.nth(0).unwrap();
if value.data.is_owned() {
// FIXME: for no-copy, need setup_op_execs to not capture tensor ptr
// mem::replace(&mut (*owner), value);
// to_replace.for_each(|t| {
// panic!("replacing");
// t.data = owner.data.view();
// });
owner.copy(&value);
} else {
owner.copy(&value);
}
} else {
println!("Unexpected input `{}`", name.as_ref());
}
}
/// Returns the graph input with name `name`, if it exists.
pub fn get_input<S: AsRef<str>>(&mut self, name: S) -> Option<&Tensor> {
self.get_input_index(name.as_ref())
.and_then(move |idx| Some(&self.tensors[idx]))
}
/// Returns the graph output with index `index`, if it exists.
pub fn get_output(&self, idx: usize) -> Option<&Tensor> {
let graph = &self.graph;
graph.heads.get(idx).and_then(|entry| {
graph
.entry_index(entry)
.map(|idx| self.tensors.get(idx))
.unwrap_or(None)
})
}
/// Returns the index for graph input with name `name`, if it exists.
pub fn get_input_index<S: AsRef<str>>(&self, name: S) -> Option<usize> {
let graph = &self.graph;
(0..graph.nodes.len())
.skip_while(|&i| graph.nodes[i].name != name.as_ref())
.nth(0)
.and_then(|i| {
if graph.arg_nodes.iter().any(|&id| id == i) {
graph.node_row_ptr.as_ref().map(|nrp| nrp[i])
} else {
None
}
})
}
}
/// Converts a string to TVM DLDataTypeCode. @see `String2TVMType` in packed_func.h
named!(
tvm_str_to_type<CompleteStr, DataType>,
do_parse!(
type_name: alpha1 >>
bits: digit1 >>
lanes: opt!(tuple!(tag!("x"), digit1)) >>
(DataType {
code: match type_name {
CompleteStr("int") => DLDataTypeCode_kDLInt,
CompleteStr("uint") => DLDataTypeCode_kDLUInt,
CompleteStr("float") => DLDataTypeCode_kDLFloat,
_ => DLDataTypeCode_kDLFloat,
} as usize,
bits: bits.parse::<u8>().unwrap() as usize,
lanes: match lanes {
Some(lanes) => lanes.1.parse::<u16>().unwrap() as usize,
None => 1,
},
})
)
);
/// Converts a bytes to String.
named!(
name<String>,
map_res!(length_bytes!(le_u64), |b: &[u8]| String::from_utf8(
b.to_vec()
))
);
/// Parses a TVMContext
named!(
tvm_ctx<&[u8], TVMContext>,
do_parse!(
device_type: le_u32 >>
device_id: le_i32 >>
(TVMContext { device_type: device_type as usize, device_id: device_id as usize })
)
);
/// Parses a DataType
named!(
data_type<&[u8], DataType>,
do_parse!(
code: le_u8 >>
bits: le_u8 >>
lanes: le_u16 >>
(DataType { code: code as usize, bits: bits as usize, lanes: lanes as usize })
)
);
/// Parses a Tensor from a TVM array file.
named!(
tensor<Tensor>,
do_parse!(
take!(8)
>> bits!(tag_bits!(u64, 64, 0))
>> ctx: tvm_ctx
>> ndim: le_u32
>> dtype: data_type
>> shape: count!(map!(le_i64, |sz| sz as i64), ndim as usize)
>> length: le_i64
>> data: take!(length)
>> (Tensor {
data: Storage::from(data),
ctx: ctx,
dtype: dtype,
size: shape.iter().product::<i64>() as usize,
shape: shape,
strides: None,
byte_offset: 0,
})
)
);
/// Parses a graph params dict from a params binary file.
named!(
parse_param_dict<HashMap<String, Tensor>>,
do_parse!(
take!(8)
>> bits!(tag_bits!(u64, 64, 0))
>> names: length_count!(le_u64, name)
>> tensors: length_count!(le_u64, tensor)
>> (HashMap::from_iter(names.into_iter().zip(tensors.into_iter())))
)
);
/// Loads a param dict saved using `nnvm.compiler.save_param_dict`.
pub fn load_param_dict(bytes: &[u8]) -> Result<HashMap<String, Tensor>> {
if let Ok((remaining_bytes, param_dict)) = parse_param_dict(bytes) {
if remaining_bytes.len() > 0 {
bail!(ErrorKind::LoadGraphParamsError("extra input".to_string()))
} else {
Ok(param_dict)
}
} else {
bail!(ErrorKind::LoadGraphParamsError(
"invalid parameters file".to_string()
))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_str_to_type() {
assert_eq!(
tvm_str_to_type(CompleteStr("float24")).unwrap().1,
DataType {
code: DLDataTypeCode_kDLFloat as usize,
bits: 24,
lanes: 1
}
);
assert_eq!(
tvm_str_to_type(CompleteStr("uint111x44")).unwrap().1,
DataType {
code: DLDataTypeCode_kDLUInt as usize,
bits: 111,
lanes: 44
}
);
}
}