440 lines
16 KiB
Rust
440 lines
16 KiB
Rust
//! Graph IR parsing implementations.
|
|
|
|
use anyhow::Result;
|
|
use serde_json::Value;
|
|
|
|
use crate::aot::{self, GraphOperation};
|
|
use crate::templates;
|
|
|
|
use super::SynthesisEngine;
|
|
|
|
impl SynthesisEngine {
|
|
/// Advanced graph IR parser supporting multiple formats (JSON, ONNX-like, custom)
|
|
pub(crate) fn parse_graph_ir(&self, graph_ir: &str) -> Result<Vec<GraphOperation>> {
|
|
tracing::debug!("Parsing graph IR: {} characters", graph_ir.len());
|
|
|
|
// Try to parse as JSON first
|
|
if let Ok(json_value) = serde_json::from_str::<Value>(graph_ir) {
|
|
return self.parse_json_ir(&json_value);
|
|
}
|
|
|
|
// Try to parse as YAML
|
|
#[cfg(feature = "yaml-parsing")]
|
|
if let Ok(yaml_value) = serde_yaml::from_str::<Value>(graph_ir) {
|
|
return self.parse_json_ir(&yaml_value); // Reuse JSON parser for YAML
|
|
}
|
|
|
|
// Try custom DSL format
|
|
if graph_ir.trim_start().starts_with("graph") {
|
|
return self.parse_custom_dsl(graph_ir);
|
|
}
|
|
|
|
// Try ONNX protobuf format detection
|
|
if graph_ir.starts_with("ir_version") || graph_ir.contains("ModelProto") {
|
|
return self.parse_onnx_like(graph_ir);
|
|
}
|
|
|
|
// Fallback: return error for unrecognized format
|
|
Err(anyhow::anyhow!(
|
|
"Unrecognized graph IR format. Supported formats: JSON, YAML, Custom DSL, ONNX-like"
|
|
))
|
|
}
|
|
|
|
/// Parse JSON/YAML format graph IR
|
|
pub(crate) fn parse_json_ir(&self, json: &Value) -> Result<Vec<GraphOperation>> {
|
|
let mut operations = Vec::new();
|
|
|
|
// Handle different JSON structures
|
|
if let Some(graph) = json.get("graph") {
|
|
return self.parse_json_graph_object(graph);
|
|
}
|
|
|
|
if let Some(nodes) = json.get("nodes").and_then(|v| v.as_array()) {
|
|
for node in nodes {
|
|
if let Some(op) = self.parse_json_node(node)? {
|
|
operations.push(op);
|
|
}
|
|
}
|
|
} else if json.is_array() {
|
|
// Direct array of operations
|
|
let nodes = json.as_array().unwrap();
|
|
for node in nodes {
|
|
if let Some(op) = self.parse_json_node(node)? {
|
|
operations.push(op);
|
|
}
|
|
}
|
|
} else {
|
|
// Single operation
|
|
if let Some(op) = self.parse_json_node(json)? {
|
|
operations.push(op);
|
|
}
|
|
}
|
|
|
|
Ok(operations)
|
|
}
|
|
|
|
/// Parse a graph object from JSON
|
|
pub(crate) fn parse_json_graph_object(&self, graph: &Value) -> Result<Vec<GraphOperation>> {
|
|
let mut operations = Vec::new();
|
|
|
|
if let Some(nodes) = graph.get("nodes").and_then(|v| v.as_array()) {
|
|
for node in nodes {
|
|
if let Some(op) = self.parse_json_node(node)? {
|
|
operations.push(op);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Handle edges/connections if present
|
|
if let Some(_edges) = graph.get("edges") {
|
|
// In a real implementation, would use edges to set up input/output connections
|
|
tracing::debug!(
|
|
"Graph contains edge information (connection parsing not fully implemented)"
|
|
);
|
|
}
|
|
|
|
Ok(operations)
|
|
}
|
|
|
|
/// Parse individual JSON node into GraphOperation
|
|
pub(crate) fn parse_json_node(&self, node: &Value) -> Result<Option<GraphOperation>> {
|
|
let id = node
|
|
.get("id")
|
|
.or_else(|| node.get("name"))
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or(&format!("op_{}", uuid::Uuid::new_v4()))
|
|
.to_string();
|
|
|
|
let op_type = node
|
|
.get("op_type")
|
|
.or_else(|| node.get("type"))
|
|
.or_else(|| node.get("operation"))
|
|
.and_then(|v| v.as_str());
|
|
|
|
let inputs = node.get("inputs").and_then(|v| v.as_array()).map_or_else(
|
|
|| vec!["input".to_string()],
|
|
|arr| {
|
|
arr.iter()
|
|
.filter_map(|v| v.as_str().map(std::string::ToString::to_string))
|
|
.collect()
|
|
},
|
|
);
|
|
|
|
let outputs = node.get("outputs").and_then(|v| v.as_array()).map_or_else(
|
|
|| vec!["output".to_string()],
|
|
|arr| {
|
|
arr.iter()
|
|
.filter_map(|v| v.as_str().map(std::string::ToString::to_string))
|
|
.collect()
|
|
},
|
|
);
|
|
|
|
if let Some(op_type_str) = op_type {
|
|
let kernel_op = self.parse_operation_type(op_type_str, node)?;
|
|
|
|
let fusable = node
|
|
.get("fusable")
|
|
.and_then(serde_json::Value::as_bool)
|
|
.unwrap_or(true);
|
|
|
|
let memory_layout = self.parse_memory_layout(node)?;
|
|
|
|
return Ok(Some(GraphOperation {
|
|
id,
|
|
operation: kernel_op,
|
|
inputs,
|
|
outputs,
|
|
memory_layout,
|
|
fusable,
|
|
}));
|
|
}
|
|
|
|
tracing::warn!("Skipping node without operation type: {}", id);
|
|
Ok(None)
|
|
}
|
|
|
|
/// Parse operation type from string and attributes
|
|
pub(crate) fn parse_operation_type(
|
|
&self,
|
|
op_type: &str,
|
|
node: &Value,
|
|
) -> Result<templates::KernelOperation> {
|
|
let attrs = node.get("attributes").or_else(|| node.get("attrs"));
|
|
|
|
match op_type.to_lowercase().as_str() {
|
|
"gemm" | "matmul" | "dense" => {
|
|
let m = self.get_int_attr(attrs, "M", 128)? as usize;
|
|
let n = self.get_int_attr(attrs, "N", 128)? as usize;
|
|
let k = self.get_int_attr(attrs, "K", 128)? as usize;
|
|
let transpose_a = self.get_bool_attr(attrs, "transpose_a", false)?;
|
|
let transpose_b = self.get_bool_attr(attrs, "transpose_b", false)?;
|
|
|
|
Ok(templates::KernelOperation::Gemm {
|
|
m: m as u32,
|
|
n: n as u32,
|
|
k: k as u32,
|
|
transpose_a,
|
|
transpose_b,
|
|
})
|
|
}
|
|
|
|
"conv2d" | "convolution" => {
|
|
let batch_size = self.get_int_attr(attrs, "batch_size", 1)? as usize;
|
|
let in_channels = self.get_int_attr(attrs, "in_channels", 64)? as usize;
|
|
let out_channels = self.get_int_attr(attrs, "out_channels", 64)? as usize;
|
|
let height = self.get_int_attr(attrs, "height", 224)? as usize;
|
|
let width = self.get_int_attr(attrs, "width", 224)? as usize;
|
|
let kernel_size = self.get_int_attr(attrs, "kernel_size", 3)? as usize;
|
|
|
|
Ok(templates::KernelOperation::Convolution {
|
|
batch_size: batch_size as u32,
|
|
in_channels: in_channels as u32,
|
|
out_channels: out_channels as u32,
|
|
height: height as u32,
|
|
width: width as u32,
|
|
kernel_size: kernel_size as u32,
|
|
})
|
|
}
|
|
|
|
"attention" | "multiheadattention" => {
|
|
let sequence_length = self.get_int_attr(attrs, "sequence_length", 512)? as usize;
|
|
let head_dim = self.get_int_attr(attrs, "head_dim", 64)? as usize;
|
|
let num_heads = self.get_int_attr(attrs, "num_heads", 8)? as usize;
|
|
|
|
Ok(templates::KernelOperation::Attention {
|
|
sequence_length: sequence_length as u32,
|
|
head_dim: head_dim as u32,
|
|
num_heads: num_heads as u32,
|
|
})
|
|
}
|
|
|
|
"add" | "sub" | "mul" | "div" | "relu" | "sigmoid" | "tanh" => {
|
|
let size = self.get_int_attr(attrs, "size", 1024)? as usize;
|
|
let elementwise_op = match op_type.to_lowercase().as_str() {
|
|
"add" => templates::ElementwiseOp::Add,
|
|
"sub" => templates::ElementwiseOp::Sub,
|
|
"mul" => templates::ElementwiseOp::Mul,
|
|
"div" => templates::ElementwiseOp::Div,
|
|
"relu" => templates::ElementwiseOp::ReLU,
|
|
"sigmoid" => templates::ElementwiseOp::Sigmoid,
|
|
"tanh" => templates::ElementwiseOp::Tanh,
|
|
_ => templates::ElementwiseOp::Add, // default
|
|
};
|
|
|
|
Ok(templates::KernelOperation::Elementwise {
|
|
operation: elementwise_op,
|
|
size: size as u32,
|
|
})
|
|
}
|
|
|
|
"sum" | "mean" | "max" | "min" => {
|
|
let input_size = self.get_int_attr(attrs, "input_size", 1024)? as usize;
|
|
let axis = self.get_int_attr(attrs, "axis", 0)? as usize;
|
|
let reduction_op = match op_type.to_lowercase().as_str() {
|
|
"sum" => templates::ReductionOp::Sum,
|
|
"mean" => templates::ReductionOp::Mean,
|
|
"max" => templates::ReductionOp::Max,
|
|
"min" => templates::ReductionOp::Min,
|
|
_ => templates::ReductionOp::Sum, // default
|
|
};
|
|
|
|
Ok(templates::KernelOperation::Reduction {
|
|
operation: reduction_op,
|
|
input_size: input_size as u32,
|
|
axis: axis as u32,
|
|
})
|
|
}
|
|
|
|
_ => {
|
|
tracing::warn!(
|
|
"Unknown operation type '{}', defaulting to elementwise add",
|
|
op_type
|
|
);
|
|
Ok(templates::KernelOperation::Elementwise {
|
|
operation: templates::ElementwiseOp::Add,
|
|
size: 1024,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Parse memory layout from node attributes
|
|
pub(crate) fn parse_memory_layout(&self, node: &Value) -> Result<aot::MemoryLayout> {
|
|
let attrs = node.get("attributes").or_else(|| node.get("attrs"));
|
|
|
|
let data_type = self.get_string_attr(attrs, "data_type", "f32")?;
|
|
let alignment = self.get_int_attr(attrs, "alignment", 32)? as usize;
|
|
|
|
Ok(aot::MemoryLayout {
|
|
dtype: data_type,
|
|
shape: vec![1], // Default shape
|
|
strides: vec![1], // Default strides
|
|
alignment: alignment as u32,
|
|
contiguous: true, // Default assumption
|
|
})
|
|
}
|
|
|
|
/// Parse custom DSL format
|
|
pub(crate) fn parse_custom_dsl(&self, dsl: &str) -> Result<Vec<GraphOperation>> {
|
|
let mut operations = Vec::new();
|
|
let lines: Vec<&str> = dsl.lines().collect();
|
|
|
|
let mut current_op: Option<GraphOperation> = None;
|
|
|
|
for line in lines {
|
|
let line = line.trim();
|
|
if line.is_empty() || line.starts_with('#') {
|
|
continue;
|
|
}
|
|
|
|
if line.starts_with("op ") {
|
|
// Save previous operation
|
|
if let Some(op) = current_op.take() {
|
|
operations.push(op);
|
|
}
|
|
|
|
// Parse operation declaration: "op operation_name type"
|
|
let parts: Vec<&str> = line.split_whitespace().collect();
|
|
if parts.len() >= 3 {
|
|
let op_id = parts[1].to_string();
|
|
let op_type = parts[2];
|
|
|
|
let operation = match op_type {
|
|
"gemm" => templates::KernelOperation::Gemm {
|
|
m: 128,
|
|
n: 128,
|
|
k: 128,
|
|
transpose_a: false,
|
|
transpose_b: false,
|
|
},
|
|
"conv2d" => templates::KernelOperation::Convolution {
|
|
batch_size: 1,
|
|
in_channels: 64,
|
|
out_channels: 64,
|
|
height: 224,
|
|
width: 224,
|
|
kernel_size: 3,
|
|
},
|
|
"add" => templates::KernelOperation::Elementwise {
|
|
operation: templates::ElementwiseOp::Add,
|
|
size: 1024,
|
|
},
|
|
_ => templates::KernelOperation::Elementwise {
|
|
operation: templates::ElementwiseOp::Add,
|
|
size: 1024,
|
|
},
|
|
};
|
|
|
|
current_op = Some(GraphOperation {
|
|
id: op_id,
|
|
operation,
|
|
inputs: vec!["input".to_string()],
|
|
outputs: vec!["output".to_string()],
|
|
memory_layout: aot::MemoryLayout::default(),
|
|
fusable: true,
|
|
});
|
|
}
|
|
} else if line.starts_with(" ") && current_op.is_some() {
|
|
// Parse operation attributes
|
|
let attr_line = line.trim();
|
|
if let Some(op) = current_op.as_mut() {
|
|
self.parse_dsl_attribute(op, attr_line)?;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Save final operation
|
|
if let Some(op) = current_op {
|
|
operations.push(op);
|
|
}
|
|
|
|
Ok(operations)
|
|
}
|
|
|
|
/// Parse DSL attribute line
|
|
pub(crate) fn parse_dsl_attribute(&self, op: &mut GraphOperation, attr_line: &str) -> Result<()> {
|
|
if let Some((key, value)) = attr_line.split_once(':') {
|
|
let key = key.trim();
|
|
let value = value.trim();
|
|
|
|
match key {
|
|
"inputs" => {
|
|
op.inputs = value.split(',').map(|s| s.trim().to_string()).collect();
|
|
}
|
|
"outputs" => {
|
|
op.outputs = value.split(',').map(|s| s.trim().to_string()).collect();
|
|
}
|
|
"fusable" => {
|
|
op.fusable = value.parse().unwrap_or(true);
|
|
}
|
|
_ => {
|
|
tracing::debug!("Unknown DSL attribute: {} = {}", key, value);
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Parse ONNX-like format (simplified)
|
|
pub(crate) fn parse_onnx_like(&self, _onnx_ir: &str) -> Result<Vec<GraphOperation>> {
|
|
// This would require a proper ONNX protobuf parser
|
|
// For now, return a placeholder implementation
|
|
tracing::warn!("ONNX parsing not fully implemented, returning default operation");
|
|
|
|
Ok(vec![GraphOperation {
|
|
id: "onnx_op".to_string(),
|
|
operation: templates::KernelOperation::Elementwise {
|
|
operation: templates::ElementwiseOp::Add,
|
|
size: 1024,
|
|
},
|
|
inputs: vec!["onnx_input".to_string()],
|
|
outputs: vec!["onnx_output".to_string()],
|
|
memory_layout: aot::MemoryLayout::default(),
|
|
fusable: true,
|
|
}])
|
|
}
|
|
|
|
// Helper methods for attribute parsing
|
|
pub(crate) fn get_int_attr(&self, attrs: Option<&Value>, key: &str, default: i64) -> Result<i64> {
|
|
if let Some(attrs) = attrs
|
|
&& let Some(value) = attrs.get(key)
|
|
{
|
|
if let Some(num) = value.as_i64() {
|
|
return Ok(num);
|
|
}
|
|
if let Some(s) = value.as_str() {
|
|
return s
|
|
.parse()
|
|
.map_err(|_| anyhow::anyhow!("Invalid integer value for {key}: {s}"));
|
|
}
|
|
}
|
|
Ok(default)
|
|
}
|
|
|
|
pub(crate) fn get_bool_attr(&self, attrs: Option<&Value>, key: &str, default: bool) -> Result<bool> {
|
|
if let Some(attrs) = attrs
|
|
&& let Some(value) = attrs.get(key)
|
|
{
|
|
if let Some(b) = value.as_bool() {
|
|
return Ok(b);
|
|
}
|
|
if let Some(s) = value.as_str() {
|
|
return s
|
|
.parse()
|
|
.map_err(|_| anyhow::anyhow!("Invalid boolean value for {key}: {s}"));
|
|
}
|
|
}
|
|
Ok(default)
|
|
}
|
|
|
|
pub(crate) fn get_string_attr(&self, attrs: Option<&Value>, key: &str, default: &str) -> Result<String> {
|
|
if let Some(attrs) = attrs
|
|
&& let Some(value) = attrs.get(key)
|
|
&& let Some(s) = value.as_str()
|
|
{
|
|
return Ok(s.to_string());
|
|
}
|
|
Ok(default.to_string())
|
|
}
|
|
}
|