osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
552 lines
17 KiB
Rust
552 lines
17 KiB
Rust
//! Tests for the synthesis engine
|
|
|
|
// Shape guards tests - active tests
|
|
use crate::aot::shape_guards::{
|
|
GuardCheckResult, GuardFailure, ShapeDim, ShapeGuard, ShapeGuardConfig, ShapeGuardManager,
|
|
ShapeSignature,
|
|
};
|
|
|
|
#[test]
|
|
fn test_shape_dim_concrete() {
|
|
let dim = ShapeDim::concrete(256);
|
|
assert!(dim.matches(256));
|
|
assert!(!dim.matches(128));
|
|
assert!(dim.is_concrete());
|
|
}
|
|
|
|
#[test]
|
|
fn test_shape_dim_symbolic() {
|
|
let dim = ShapeDim::symbolic("batch_size");
|
|
assert!(dim.matches(1));
|
|
assert!(dim.matches(256));
|
|
assert!(dim.matches(10000));
|
|
assert!(!dim.is_concrete());
|
|
assert_eq!(dim.symbolic_name(), Some("batch_size"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_shape_dim_bounded() {
|
|
let dim = ShapeDim::bounded("seq_len", 1, 512);
|
|
assert!(dim.matches(1));
|
|
assert!(dim.matches(256));
|
|
assert!(dim.matches(512));
|
|
assert!(!dim.matches(0));
|
|
assert!(!dim.matches(513));
|
|
}
|
|
|
|
#[test]
|
|
fn test_shape_guard_check_pass() {
|
|
let guard = ShapeGuard::concrete("input_0", &[32, 256, 256]);
|
|
|
|
match guard.check(&[32, 256, 256], Some("fp32")) {
|
|
GuardCheckResult::Passed => {}
|
|
GuardCheckResult::Failed(f) => panic!("Guard should pass: {:?}", f),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_shape_guard_check_rank_mismatch() {
|
|
let guard = ShapeGuard::concrete("input_0", &[32, 256, 256]);
|
|
|
|
match guard.check(&[32, 256], None) {
|
|
GuardCheckResult::Failed(GuardFailure::RankMismatch { .. }) => {}
|
|
other => panic!("Expected RankMismatch, got {:?}", other),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_shape_guard_check_dim_mismatch() {
|
|
let guard = ShapeGuard::concrete("input_0", &[32, 256, 256]);
|
|
|
|
match guard.check(&[64, 256, 256], None) {
|
|
GuardCheckResult::Failed(GuardFailure::DimensionMismatch { dim: 0, .. }) => {}
|
|
other => panic!("Expected DimensionMismatch at dim 0, got {:?}", other),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_shape_guard_symbolic_bindings() {
|
|
let guard = ShapeGuard::new(
|
|
"input_0",
|
|
vec![
|
|
ShapeDim::symbolic("batch"),
|
|
ShapeDim::concrete(256),
|
|
ShapeDim::symbolic("seq_len"),
|
|
],
|
|
);
|
|
|
|
let bindings = guard.extract_bindings(&[32, 256, 512]);
|
|
assert_eq!(bindings.get("batch"), Some(&32));
|
|
assert_eq!(bindings.get("seq_len"), Some(&512));
|
|
assert_eq!(bindings.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shape_signature_hash() {
|
|
let inputs1 = vec![("input_0".to_string(), vec![32, 256], "fp32".to_string())];
|
|
let inputs2 = vec![("input_0".to_string(), vec![32, 256], "fp32".to_string())];
|
|
let inputs3 = vec![("input_0".to_string(), vec![64, 256], "fp32".to_string())];
|
|
|
|
let sig1 = ShapeSignature::from_inputs(&inputs1, 12345);
|
|
let sig2 = ShapeSignature::from_inputs(&inputs2, 12345);
|
|
let sig3 = ShapeSignature::from_inputs(&inputs3, 12345);
|
|
|
|
assert_eq!(sig1.to_cache_key(), sig2.to_cache_key());
|
|
assert_ne!(sig1.to_cache_key(), sig3.to_cache_key());
|
|
}
|
|
|
|
#[test]
|
|
fn test_shape_guard_manager_stats() {
|
|
let manager = ShapeGuardManager::default_manager();
|
|
assert_eq!(manager.stats().total_checks, 0);
|
|
assert_eq!(manager.stats().cache_hits, 0);
|
|
}
|
|
|
|
#[cfg(feature = "disabled_tests")]
|
|
mod integration_tests {
|
|
use super::*;
|
|
use crate::{
|
|
SynthesisEngine,
|
|
aot::{CompiledGraph, ExecutionPlan},
|
|
templates,
|
|
};
|
|
|
|
#[tokio::test]
|
|
async fn test_synthesis_engine_creation() {
|
|
let result = SynthesisEngine::new("sm_120");
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_synthesis_engine_initialization() {
|
|
let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
let result = engine.initialize().await;
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_json_ir_parsing() {
|
|
let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
|
|
// Test simple JSON operation
|
|
let json_ir = r#"{
|
|
"id": "test_op",
|
|
"op_type": "gemm",
|
|
"attributes": {
|
|
"M": 256,
|
|
"N": 256,
|
|
"K": 256,
|
|
"transpose_a": false,
|
|
"transpose_b": true
|
|
},
|
|
"inputs": ["input_a", "input_b"],
|
|
"outputs": ["output_c"]
|
|
}"#;
|
|
|
|
let operations = engine
|
|
.parse_graph_ir(json_ir)
|
|
.expect("Should parse JSON successfully");
|
|
assert_eq!(operations.len(), 1);
|
|
|
|
let op = &operations[0];
|
|
assert_eq!(op.id, "test_op");
|
|
assert_eq!(op.inputs, vec!["input_a", "input_b"]);
|
|
assert_eq!(op.outputs, vec!["output_c"]);
|
|
|
|
if let templates::KernelOperation::Gemm {
|
|
m,
|
|
n,
|
|
k,
|
|
transpose_a,
|
|
transpose_b,
|
|
} = op.operation
|
|
{
|
|
assert_eq!(m, 256);
|
|
assert_eq!(n, 256);
|
|
assert_eq!(k, 256);
|
|
assert!(!transpose_a);
|
|
assert!(transpose_b);
|
|
} else {
|
|
panic!("Expected GEMM operation");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_json_array_ir_parsing() {
|
|
let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
|
|
// Test array of operations
|
|
let json_ir = r#"[
|
|
{
|
|
"id": "conv_op",
|
|
"op_type": "conv2d",
|
|
"attributes": {
|
|
"batch_size": 1,
|
|
"in_channels": 64,
|
|
"out_channels": 128,
|
|
"height": 224,
|
|
"width": 224,
|
|
"kernel_size": 3
|
|
}
|
|
},
|
|
{
|
|
"id": "relu_op",
|
|
"op_type": "relu",
|
|
"attributes": {
|
|
"size": 200704
|
|
}
|
|
}
|
|
]"#;
|
|
|
|
let operations = engine
|
|
.parse_graph_ir(json_ir)
|
|
.expect("Should parse JSON array successfully");
|
|
assert_eq!(operations.len(), 2);
|
|
|
|
// Check convolution
|
|
if let templates::KernelOperation::Convolution {
|
|
batch_size,
|
|
in_channels,
|
|
out_channels,
|
|
height,
|
|
width,
|
|
kernel_size,
|
|
} = operations[0].operation
|
|
{
|
|
assert_eq!(batch_size, 1);
|
|
assert_eq!(in_channels, 64);
|
|
assert_eq!(out_channels, 128);
|
|
assert_eq!(height, 224);
|
|
assert_eq!(width, 224);
|
|
assert_eq!(kernel_size, 3);
|
|
} else {
|
|
panic!("Expected Convolution operation");
|
|
}
|
|
|
|
// Check ReLU
|
|
if let templates::KernelOperation::Elementwise { operation, size } = operations[1].operation
|
|
{
|
|
assert_eq!(size, 200704);
|
|
if let templates::ElementwiseOp::ReLU = operation {
|
|
// Correct
|
|
} else {
|
|
panic!("Expected ReLU elementwise operation");
|
|
}
|
|
} else {
|
|
panic!("Expected Elementwise operation");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_custom_dsl_parsing() {
|
|
let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
|
|
let dsl_ir = r#"
|
|
# Test DSL format
|
|
op matrix_mult gemm
|
|
inputs: A, B
|
|
outputs: C
|
|
fusable: true
|
|
|
|
op activation add
|
|
inputs: C, bias
|
|
outputs: result
|
|
fusable: true
|
|
"#;
|
|
|
|
let operations = engine
|
|
.parse_graph_ir(dsl_ir)
|
|
.expect("Should parse DSL successfully");
|
|
assert_eq!(operations.len(), 2);
|
|
|
|
// Check first operation
|
|
let gemm_op = &operations[0];
|
|
assert_eq!(gemm_op.id, "matrix_mult");
|
|
assert_eq!(gemm_op.inputs, vec!["A", "B"]);
|
|
assert_eq!(gemm_op.outputs, vec!["C"]);
|
|
assert!(gemm_op.fusable);
|
|
|
|
// Check second operation
|
|
let add_op = &operations[1];
|
|
assert_eq!(add_op.id, "activation");
|
|
assert_eq!(add_op.inputs, vec!["C", "bias"]);
|
|
assert_eq!(add_op.outputs, vec!["result"]);
|
|
assert!(add_op.fusable);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_graph_object_parsing() {
|
|
let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
|
|
let graph_json = r#"{
|
|
"graph": {
|
|
"nodes": [
|
|
{
|
|
"id": "attention_op",
|
|
"op_type": "attention",
|
|
"attributes": {
|
|
"sequence_length": 1024,
|
|
"head_dim": 64,
|
|
"num_heads": 16
|
|
}
|
|
}
|
|
],
|
|
"edges": [
|
|
{"from": "input", "to": "attention_op"}
|
|
]
|
|
}
|
|
}"#;
|
|
|
|
let operations = engine
|
|
.parse_graph_ir(graph_json)
|
|
.expect("Should parse graph JSON successfully");
|
|
assert_eq!(operations.len(), 1);
|
|
|
|
if let templates::KernelOperation::Attention {
|
|
sequence_length,
|
|
head_dim,
|
|
num_heads,
|
|
} = operations[0].operation
|
|
{
|
|
assert_eq!(sequence_length, 1024);
|
|
assert_eq!(head_dim, 64);
|
|
assert_eq!(num_heads, 16);
|
|
} else {
|
|
panic!("Expected Attention operation");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_flops_calculation() {
|
|
let engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
|
|
// Test GEMM FLOPS
|
|
let gemm_op = templates::KernelOperation::Gemm {
|
|
m: 128,
|
|
n: 128,
|
|
k: 128,
|
|
transpose_a: false,
|
|
transpose_b: false,
|
|
};
|
|
let flops = engine.calculate_theoretical_flops(&gemm_op);
|
|
let expected = 2.0 * 128.0 * 128.0 * 128.0; // 2 * M * N * K
|
|
assert!(
|
|
(flops - expected).abs() < 1.0,
|
|
"GEMM FLOPS calculation incorrect"
|
|
);
|
|
|
|
// Test Convolution FLOPS
|
|
let conv_op = templates::KernelOperation::Convolution {
|
|
batch_size: 1,
|
|
in_channels: 64,
|
|
out_channels: 128,
|
|
height: 224,
|
|
width: 224,
|
|
kernel_size: 3,
|
|
};
|
|
let conv_flops = engine.calculate_theoretical_flops(&conv_op);
|
|
let expected_conv = 2.0 * 1.0 * 128.0 * 224.0 * 224.0 * 64.0 * 9.0; // 2 * B * OC * H * W * IC * K^2
|
|
assert!(
|
|
(conv_flops - expected_conv).abs() < 1.0,
|
|
"Convolution FLOPS calculation incorrect"
|
|
);
|
|
|
|
// Test Elementwise FLOPS
|
|
let elem_op = templates::KernelOperation::Elementwise {
|
|
operation: templates::ElementwiseOp::Add,
|
|
size: 1024,
|
|
};
|
|
let elem_flops = engine.calculate_theoretical_flops(&elem_op);
|
|
assert_eq!(elem_flops, 1024.0, "Elementwise FLOPS should equal size");
|
|
}
|
|
|
|
#[test]
|
|
fn test_memory_usage_estimation() {
|
|
let engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
|
|
// Test GEMM memory usage
|
|
let gemm_op = templates::KernelOperation::Gemm {
|
|
m: 128,
|
|
n: 128,
|
|
k: 128,
|
|
transpose_a: false,
|
|
transpose_b: false,
|
|
};
|
|
let memory = engine.estimate_memory_usage(&gemm_op);
|
|
let expected = 4.0 * ((128 * 128) + (128 * 128) + (128 * 128)) as f64; // A + B + C in FP32
|
|
assert!(
|
|
(memory - expected).abs() < 1.0,
|
|
"GEMM memory estimation incorrect"
|
|
);
|
|
|
|
// Test Elementwise memory usage
|
|
let elem_op = templates::KernelOperation::Elementwise {
|
|
operation: templates::ElementwiseOp::Add,
|
|
size: 1024,
|
|
};
|
|
let elem_memory = engine.estimate_memory_usage(&elem_op);
|
|
assert_eq!(
|
|
elem_memory,
|
|
8.0 * 1024.0,
|
|
"Elementwise memory should be 2 * size * sizeof(f32)"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_parameter_space_creation() {
|
|
let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
engine.initialize().await.expect("Should initialize");
|
|
|
|
let template = templates::KernelTemplate::new("test_template".to_string());
|
|
let gemm_op = templates::KernelOperation::Gemm {
|
|
m: 256,
|
|
n: 256,
|
|
k: 256,
|
|
transpose_a: false,
|
|
transpose_b: false,
|
|
};
|
|
|
|
let param_space = engine
|
|
.create_parameter_space(&template, &gemm_op)
|
|
.expect("Should create parameter space");
|
|
|
|
// Check that GEMM-specific parameters are included
|
|
assert!(param_space.dimensions.contains_key("tile_m"));
|
|
assert!(param_space.dimensions.contains_key("tile_n"));
|
|
assert!(param_space.dimensions.contains_key("tile_k"));
|
|
assert!(param_space.dimensions.contains_key("block_size_x"));
|
|
assert!(param_space.dimensions.contains_key("block_size_y"));
|
|
|
|
// Check constraints
|
|
assert!(!param_space.constraints.is_empty());
|
|
assert!(
|
|
param_space
|
|
.constraints
|
|
.iter()
|
|
.any(|c| c.contains("block_size_x * block_size_y"))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parameter_conversion() {
|
|
let engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
|
|
let mut int_params = std::collections::HashMap::new();
|
|
int_params.insert("block_size_x".to_string(), 128);
|
|
int_params.insert("block_size_y".to_string(), 8);
|
|
int_params.insert("tile_m".to_string(), 64);
|
|
|
|
let template_params = engine
|
|
.convert_to_template_params(&int_params)
|
|
.expect("Should convert parameters");
|
|
|
|
assert_eq!(template_params["block_size_x"], "128");
|
|
assert_eq!(template_params["block_size_y"], "8");
|
|
assert_eq!(template_params["tile_m"], "64");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_invalid_ir_formats() {
|
|
let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
|
|
// Test invalid JSON
|
|
let invalid_json = r#"{ invalid json }"#;
|
|
let result = engine.parse_graph_ir(invalid_json);
|
|
assert!(result.is_err(), "Should reject invalid JSON");
|
|
|
|
// Test unknown format
|
|
let unknown_format = "this is not a known IR format";
|
|
let result = engine.parse_graph_ir(unknown_format);
|
|
assert!(result.is_err(), "Should reject unknown format");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compiled_kernel_structure() {
|
|
use crate::CompiledKernel;
|
|
|
|
let compiled_graph = CompiledGraph {
|
|
kernel_sequence: vec![],
|
|
memory_plan: std::collections::HashMap::new(),
|
|
execution_plan: ExecutionPlan {
|
|
kernel_launches: vec![],
|
|
sync_points: vec![],
|
|
},
|
|
};
|
|
|
|
let kernel = CompiledKernel {
|
|
code: "test kernel code".to_string(),
|
|
compiled_graph,
|
|
};
|
|
|
|
assert_eq!(kernel.code, "test kernel code");
|
|
assert!(kernel.compiled_graph.kernel_sequence.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tuning_config_creation() {
|
|
let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
engine.initialize().await.expect("Should initialize");
|
|
|
|
// Test GEMM config (should have more iterations)
|
|
let gemm_op = templates::KernelOperation::Gemm {
|
|
m: 128,
|
|
n: 128,
|
|
k: 128,
|
|
transpose_a: false,
|
|
transpose_b: false,
|
|
};
|
|
let gemm_config = engine
|
|
.create_tuning_config(&gemm_op)
|
|
.expect("Should create GEMM config");
|
|
assert_eq!(gemm_config.max_iterations, 100);
|
|
|
|
// Test Attention config
|
|
let attention_op = templates::KernelOperation::Attention {
|
|
sequence_length: 512,
|
|
head_dim: 64,
|
|
num_heads: 8,
|
|
};
|
|
let attention_config = engine
|
|
.create_tuning_config(&attention_op)
|
|
.expect("Should create Attention config");
|
|
assert_eq!(attention_config.max_iterations, 80);
|
|
|
|
// Test default config for other operations
|
|
let elem_op = templates::KernelOperation::Elementwise {
|
|
operation: templates::ElementwiseOp::Add,
|
|
size: 1024,
|
|
};
|
|
let elem_config = engine
|
|
.create_tuning_config(&elem_op)
|
|
.expect("Should create default config");
|
|
assert_eq!(elem_config.max_iterations, 50);
|
|
}
|
|
|
|
#[test]
|
|
fn test_execution_time_estimation() {
|
|
let mut engine = SynthesisEngine::new("sm_120").expect("Failed to create engine");
|
|
|
|
let mock_graph = CompiledGraph {
|
|
kernel_sequence: vec![],
|
|
memory_plan: std::collections::HashMap::new(),
|
|
execution_plan: ExecutionPlan {
|
|
kernel_launches: vec![],
|
|
sync_points: vec![],
|
|
},
|
|
};
|
|
|
|
let gemm_op = templates::KernelOperation::Gemm {
|
|
m: 1024,
|
|
n: 1024,
|
|
k: 1024,
|
|
transpose_a: false,
|
|
transpose_b: false,
|
|
};
|
|
let exec_time = engine.estimate_execution_time(&gemm_op, &mock_graph);
|
|
|
|
// Should be reasonable execution time (not zero, not too large)
|
|
assert!(exec_time.as_secs_f64() > 0.0);
|
|
assert!(exec_time.as_secs_f64() < 10.0); // Should complete within 10 seconds for reasonable sizes
|
|
}
|
|
}
|