Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,439 @@
|
||||
//! 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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user