From 15e35d92957c86769a0ecbd6d086f133bfbcf721 Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Fri, 2 Jul 2021 23:51:32 -0700 Subject: [PATCH 01/11] codegen macro PR intial commit (draft) --- src/codegen.rs | 79 +++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 78 insertions(+), 1 deletion(-) diff --git a/src/codegen.rs b/src/codegen.rs index 61ccae4421..61e87b1695 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -1938,7 +1938,55 @@ mod tests { } relay_outputs } - + macro_rules! codegen_test { + // pattern where shape/map is present + //assorted variables as tuple? + //code_expr is the glenside expression, (needs to be whole thing, not a format string, variable, etc.) + ($test_name: ident, $shape_vec: expr, $code_expr: expr, $code_tup: expr) => { + #[test] + fn $test_name() { + let mut map = HashMap::default(); + let shape_vec = $shape_vec; + if (shape_vec.len() == 1) { + map.insert("t".to_string(), shape_vec[0].0.clone()); + } else { + for n in 0..shape_vec.len() { + map.insert(format!("t{}", n).to_string(), shape_vec[n].0.clone()); + } + } + let expr = RecExpr::from_str($code_expr.as_str()).unwrap(); + let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); + let id = egraph.add_expr(&expr); + let code = codegen( + &egraph, + id, + &$code_tup.0, + $code_tup.1, + "", + &$code_tup.2, + &generate_worklist_for_codegen(&egraph, id), + true, + ); + } + }; + } + codegen_test!( + tranpose, + vec![( + vec![1, 20, 300, 3], + ndarray::ArrayD::from_shape_vec( + vec![1, 20, 300, 3].clone(), + (0..vec![1, 20, 300, 3].iter().product::()).collect(), + ) + .unwrap() + )], + format!( + " +(access-transpose (access-tensor t) (list {}))", + vec![3, 1, 0, 2].iter().map(|x| x.to_string()).join(" ") + ), + (HashMap::default(), "transpose", vec!["t"]) + ); #[test] fn transpose() { let shape = vec![1, 20, 300, 3]; @@ -2438,6 +2486,35 @@ int main() {{ #[test] fn slice() { + + // our_macro!( + // { + // let slice_axis = 2; + // let low = 5; + // let high = 83; + // let mut slices = Vec::from_iter( + // std::iter::repeat(SliceOrIndex::Slice { + // start: 0, + // end: None, + // step: 1, + // }) + // .take(shape.len()), + // ); + // slices[slice_axis] = SliceOrIndex::Slice { + // start: low, + // end: Some(high), + // step: 1, + // }; + // let sliced = input.slice( + // &SliceInfo::, ndarray::IxDyn>::new(slices) + // .unwrap() + // .as_ref(), + // ); + + // sliced + // }, + // ) + let shape = vec![32, 7, 100, 3]; let slice_axis = 2; let low = 5; From 919e7245a7e383f9e03fb8fe3562eb1e0ec71413 Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Mon, 12 Jul 2021 20:28:38 -0700 Subject: [PATCH 02/11] Initial test rewrites with macro (tranpose, concat, pad) --- src/codegen.rs | 583 +++++++++++++++++-------------------------------- 1 file changed, 196 insertions(+), 387 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index 61e87b1695..fae696cd41 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -1939,281 +1939,201 @@ mod tests { relay_outputs } macro_rules! codegen_test { - // pattern where shape/map is present - //assorted variables as tuple? - //code_expr is the glenside expression, (needs to be whole thing, not a format string, variable, etc.) - ($test_name: ident, $shape_vec: expr, $code_expr: expr, $code_tup: expr) => { - #[test] - fn $test_name() { - let mut map = HashMap::default(); - let shape_vec = $shape_vec; - if (shape_vec.len() == 1) { - map.insert("t".to_string(), shape_vec[0].0.clone()); - } else { - for n in 0..shape_vec.len() { - map.insert(format!("t{}", n).to_string(), shape_vec[n].0.clone()); - } + ($input_vec: expr, $code_expr: expr, $code_tup: expr, $ground_truth : expr, $c_code : expr) => { + let mut map: HashMap> = HashMap::default(); + let shape_vec = $input_vec; + let mut names = Vec::new(); + let expr: RecExpr = RecExpr::from_str($code_expr.as_str()).unwrap(); + if (shape_vec.len() == 1) { + let name = shape_vec[0].0; + map.insert(name.to_string(), shape_vec[0].1.shape().to_vec()); + names.push(name); + } else { + for n in 0..(shape_vec.len()) { + let name = shape_vec[n].0.clone(); + let shape = shape_vec[n].1.shape().to_vec().clone(); + + map.insert(name.to_string(), shape); + names.push(name); } - let expr = RecExpr::from_str($code_expr.as_str()).unwrap(); - let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); - let id = egraph.add_expr(&expr); - let code = codegen( - &egraph, - id, - &$code_tup.0, - $code_tup.1, - "", - &$code_tup.2, - &generate_worklist_for_codegen(&egraph, id), - true, - ); } + let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); + let id = egraph.add_expr(&expr); + let code = codegen( + &egraph, + id, + &$code_tup.0, + $code_tup.1, + "", + &names, + &generate_worklist_for_codegen(&egraph, id), + true, + ); + let ground_truth = $ground_truth; + let c_func = format!( + " + int main() {{ + {}(out{}); + for (int i = 0; i < {}; i++) {{ + assert(((float*)ground_truth)[i] == ((float*)out)[i]); + }} + }} + ", + $code_tup.1, + names + .iter() + .fold(String::new(), |acc, &arg| acc + "," + arg), + ground_truth.iter().product::() + ); + let declarations = format!( + " + #include + {} + {} + {} + {} + {}", + $c_code, + shape_vec.iter().fold(String::new(), |acc, arg| { + format!( + "{}{}{}", + acc, + c_assignment_string("", arg.0, DType::Fp32, &arg.1.view()), + "\n" + ) + }), + c_assignment_string("", "ground_truth", DType::Fp32, &ground_truth.view()), + c_assignment_string( + "", + "out", + DType::Fp32, + &ndarray::ArrayD::::zeros(ground_truth.shape()).view() + ), + code + ); + + let main_code = format!( + " + {} + {}", + declarations, c_func + ); + let main_c_filepath = std::env::temp_dir().with_file_name(format!( + "{}-test-{}.c", + $code_tup.1, + std::time::SystemTime::now().elapsed().unwrap().as_nanos() + )); + let binary_filepath = std::env::temp_dir().with_file_name(format!( + "{}-test-{}", + $code_tup.1, + std::time::SystemTime::now().elapsed().unwrap().as_nanos() + )); + + File::create(&main_c_filepath) + .unwrap() + .write_all(main_code.as_bytes()) + .unwrap(); + + let result = Command::new("gcc") + .arg("-Werror") + .arg("-g") + .arg("-o") + .arg(&binary_filepath) + .arg(&main_c_filepath) + .output() + .unwrap(); + + assert!( + result.status.success(), + "{}", + std::str::from_utf8(result.stderr.as_slice()) + .expect("Could not convert stderr to UTF8") + ); + + let result = Command::new(&binary_filepath).output().unwrap(); + + assert!( + result.status.success(), + "{}", + std::str::from_utf8(result.stderr.as_slice()) + .expect("Could not convert stderr to UTF8") + ); }; } - codegen_test!( - tranpose, - vec![( - vec![1, 20, 300, 3], + #[test] + fn tranpose() { + let permutation = vec![3, 1, 0, 2]; + let input_list = vec![( + "t", ndarray::ArrayD::from_shape_vec( vec![1, 20, 300, 3].clone(), (0..vec![1, 20, 300, 3].iter().product::()).collect(), ) - .unwrap() - )], - format!( - " -(access-transpose (access-tensor t) (list {}))", - vec![3, 1, 0, 2].iter().map(|x| x.to_string()).join(" ") - ), - (HashMap::default(), "transpose", vec!["t"]) - ); - #[test] - fn transpose() { - let shape = vec![1, 20, 300, 3]; - let permutation = vec![3, 1, 0, 2]; - let input = ndarray::ArrayD::from_shape_vec( - shape.clone(), - (0..shape.iter().product::()).collect(), - ) - .unwrap(); - let input_transposed = input.clone().permuted_axes(permutation.clone()); - - let expr = RecExpr::from_str( + .unwrap(), + )]; + let input_transposed = input_list[0].1.clone().permuted_axes(permutation.clone()); + codegen_test!( + input_list.clone(), format!( - " -(access-transpose (access-tensor t) (list {}))", + "(access-transpose (access-tensor t) (list {}))", permutation.iter().map(|x| x.to_string()).join(" ") - ) - .as_str(), - ) - .unwrap(); - - let mut map = HashMap::default(); - map.insert("t".to_string(), shape.clone()); - - let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); - let id = egraph.add_expr(&expr); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "transpose", - "", - &vec!["t"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include - -{} -{} -{} -{} - -int main() {{ - transpose(out, a); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)a_t)[i] == ((float*)out)[i]); - }} -}} -", - c_assignment_string("", "a", DType::Fp32, &input.view()), - c_assignment_string("", "a_t", DType::Fp32, &input_transposed.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(input_transposed.shape()).view() ), - code, - shape.iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "transpose-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "transpose-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") + ( + { + let hw_map: HashMap = HashMap::default(); + hw_map + }, + "tranpose" + ), + input_transposed, + "" ); } #[test] fn concat() { - let shape0 = vec![2, 10, 50, 3]; - let shape1 = vec![2, 3, 50, 3]; + let input_list = vec![ + ( + "t0", + ndarray::ArrayD::from_shape_vec( + vec![2, 10, 50, 3].clone(), + (0..vec![2, 10, 50, 3].iter().product::()).collect(), + ) + .unwrap(), + ), + ( + "t1", + ndarray::ArrayD::from_shape_vec( + vec![2, 3, 50, 3].clone(), + (0..vec![2, 3, 50, 3].iter().product::()).collect(), + ) + .unwrap(), + ), + ]; let concat_axis = 1; - - let input0 = ndarray::ArrayD::from_shape_vec( - shape0.clone(), - (0..shape0.iter().product::()).collect(), - ) - .unwrap(); - let input1 = ndarray::ArrayD::from_shape_vec( - shape1.clone(), - (0..shape1.iter().product::()).collect(), + let concatted = ndarray::stack( + ndarray::Axis(concat_axis), + &[ + input_list[0].1.clone().view(), + input_list[1].1.clone().view(), + ], ) .unwrap(); - let concatted = - ndarray::stack(ndarray::Axis(concat_axis), &[input0.view(), input1.view()]).unwrap(); - - let expr = RecExpr::from_str( + codegen_test!( + input_list.clone(), format!( - " -(access-concatenate (access-tensor t0) (access-tensor t1) {})", + "(access-concatenate (access-tensor t0) (access-tensor t1) {})", concat_axis - ) - .as_str(), - ) - .unwrap(); - - let mut map = HashMap::default(); - map.insert("t0".to_string(), shape0.clone()); - map.insert("t1".to_string(), shape1.clone()); - - let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); - let id = egraph.add_expr(&expr); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "concatenate", - "", - &vec!["t0", "t1"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include - -{} -{} -{} -{} -{} - -int main() {{ - concatenate(out, t0, t1); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)a_t)[i] == ((float*)out)[i]); - }} -}} -", - c_assignment_string("", "t0", DType::Fp32, &input0.view()), - c_assignment_string("", "t1", DType::Fp32, &input1.view()), - c_assignment_string("", "a_t", DType::Fp32, &concatted.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(concatted.shape()).view() ), - code, - concatted.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "concatenate-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "concatenate-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") + ( + { + let hw_map: HashMap = HashMap::default(); + hw_map + }, + "concatenate" + ), + concatted, + "" ); } @@ -2354,167 +2274,56 @@ int main() {{ .expect("Could not convert stderr to UTF8") ); } - #[test] fn pad() { - let shape = vec![10, 20, 3, 45]; + let input_list = vec![( + "t", + ndarray::ArrayD::from_shape_vec( + vec![10, 20, 3, 45].clone(), + (0..vec![10, 20, 3, 45].iter().product::()).collect(), + ) + .unwrap(), + )]; let pad_axis = 2; let pad_before = 2; let pad_after = 3; let pad_type = PadType::ZeroPadding; - let mut pad_before_shape = shape.clone(); + let mut pad_before_shape = input_list[0].1.shape().to_vec().clone(); pad_before_shape[pad_axis] = pad_before; - let mut pad_after_shape = shape.clone(); + let mut pad_after_shape = input_list[0].1.shape().to_vec().clone(); pad_after_shape[pad_axis] = pad_after; - - let input = ndarray::ArrayD::from_shape_vec( - shape.clone(), - (0..shape.iter().product::()).collect(), - ) - .unwrap(); - assert!(pad_type == PadType::ZeroPadding); let padded = ndarray::stack( ndarray::Axis(pad_axis), &[ ndarray::ArrayD::zeros(pad_before_shape).view(), - input.view(), + input_list[0].1.view(), ndarray::ArrayD::zeros(pad_after_shape).view(), ], ) .unwrap(); - - let expr = RecExpr::from_str( + codegen_test!( + input_list.clone(), format!( - " -(access-pad (access-tensor t) {} {} {} {})", + "(access-pad (access-tensor t) {} {} {} {})", pad_type, pad_axis, pad_before, pad_after - ) - .as_str(), - ) - .unwrap(); - - let mut map = HashMap::default(); - map.insert("t".to_string(), shape.clone()); - - let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); - let id = egraph.add_expr(&expr); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "pad", - "", - &vec!["t"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include - -{} -{} -{} -{} - -int main() {{ - pad(out, a); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)a_pad)[i] == ((float*)out)[i]); - }} -}} -", - c_assignment_string("", "a", DType::Fp32, &input.view()), - c_assignment_string("", "a_pad", DType::Fp32, &padded.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(padded.shape()).view() ), - code, - shape.iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "pad-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "pad-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") + ( + { + let hw_map: HashMap = HashMap::default(); + hw_map + }, + "concatenate" + ), + padded, + "" ); } #[test] fn slice() { - - // our_macro!( - // { - // let slice_axis = 2; - // let low = 5; - // let high = 83; - // let mut slices = Vec::from_iter( - // std::iter::repeat(SliceOrIndex::Slice { - // start: 0, - // end: None, - // step: 1, - // }) - // .take(shape.len()), - // ); - // slices[slice_axis] = SliceOrIndex::Slice { - // start: low, - // end: Some(high), - // step: 1, - // }; - // let sliced = input.slice( - // &SliceInfo::, ndarray::IxDyn>::new(slices) - // .unwrap() - // .as_ref(), - // ); - - // sliced - // }, - // ) - let shape = vec![32, 7, 100, 3]; let slice_axis = 2; let low = 5; From f04b779464daa6ce5dc22b0f1d8c3c67f6ae06ac Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Wed, 14 Jul 2021 23:44:38 -0700 Subject: [PATCH 03/11] refactor macro --- src/codegen.rs | 74 ++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index fae696cd41..a53cc8df5c 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -1939,19 +1939,19 @@ mod tests { relay_outputs } macro_rules! codegen_test { - ($input_vec: expr, $code_expr: expr, $code_tup: expr, $ground_truth : expr, $c_code : expr) => { + ($env: expr, $code_expr: expr, $code_tup: expr, $ground_truth : expr, $c_code : expr) => { let mut map: HashMap> = HashMap::default(); - let shape_vec = $input_vec; + let env = $env; let mut names = Vec::new(); let expr: RecExpr = RecExpr::from_str($code_expr.as_str()).unwrap(); - if (shape_vec.len() == 1) { - let name = shape_vec[0].0; - map.insert(name.to_string(), shape_vec[0].1.shape().to_vec()); + if (env.len() == 1) { + let name = env[0].0; + map.insert(name.to_string(), env[0].1.shape().to_vec()); names.push(name); } else { - for n in 0..(shape_vec.len()) { - let name = shape_vec[n].0.clone(); - let shape = shape_vec[n].1.shape().to_vec().clone(); + for n in 0..(env.len()) { + let name = env[n].0.clone(); + let shape = env[n].1.shape().to_vec().clone(); map.insert(name.to_string(), shape); names.push(name); @@ -1963,7 +1963,7 @@ mod tests { &egraph, id, &$code_tup.0, - $code_tup.1, + "compiled_function", "", &names, &generate_worklist_for_codegen(&egraph, id), @@ -1983,7 +1983,7 @@ mod tests { names .iter() .fold(String::new(), |acc, &arg| acc + "," + arg), - ground_truth.iter().product::() + ground_truth.shape().iter().product::() ); let declarations = format!( " @@ -1994,7 +1994,7 @@ mod tests { {} {}", $c_code, - shape_vec.iter().fold(String::new(), |acc, arg| { + env.iter().fold(String::new(), |acc, arg| { format!( "{}{}{}", acc, @@ -2321,7 +2321,57 @@ int main() {{ "" ); } - + #[test] + fn test() { + let input_list = vec![ + ( + "t", + ndarray::ArrayD::from_shape_vec( + vec![32, 7, 100, 3].clone(), + (0..vec![32, 7, 100, 3].iter().product::()).collect(), + ) + .unwrap() + ) + ]; + let slice_axis = 2; + let low = 5; + let high = 83; + let mut slices = Vec::from_iter( + std::iter::repeat(SliceOrIndex::Slice { + start: 0, + end: None, + step: 1, + }) + .take(input_list[0].1.shape().to_vec().len()), + ); + slices[slice_axis] = SliceOrIndex::Slice { + start: low, + end: Some(high), + step: 1, + }; + let sliced = input_list[0].1.slice( + &SliceInfo::, ndarray::IxDyn>::new(slices) + .unwrap() + .as_ref(), + ); + codegen_test!( + input_list.clone(), + format!( + " +(access-slice (access-tensor t) {} {} {})", + slice_axis, low, high + ), + ( + { + let hw_map: HashMap = HashMap::default(); + hw_map + }, + "slice" + ), + sliced, + "" + ); + } #[test] fn slice() { let shape = vec![32, 7, 100, 3]; From 91e2aac04ed116c5d71dcb412f11eb9c150aa301 Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Fri, 16 Jul 2021 13:24:22 -0700 Subject: [PATCH 04/11] temp changes --- src/codegen.rs | 170 ++----------------------------------------------- 1 file changed, 7 insertions(+), 163 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index a53cc8df5c..6bd6ed7fba 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -1939,7 +1939,7 @@ mod tests { relay_outputs } macro_rules! codegen_test { - ($env: expr, $code_expr: expr, $code_tup: expr, $ground_truth : expr, $c_code : expr) => { + ($env: expr, $code_expr: expr, $ground_truth : expr, $c_code : expr) => { let mut map: HashMap> = HashMap::default(); let env = $env; let mut names = Vec::new(); @@ -1962,7 +1962,7 @@ mod tests { let code = codegen( &egraph, id, - &$code_tup.0, + &HashMap::default(), "compiled_function", "", &names, @@ -1979,7 +1979,7 @@ mod tests { }} }} ", - $code_tup.1, + "compiled_function", names .iter() .fold(String::new(), |acc, &arg| acc + "," + arg), @@ -2020,12 +2020,12 @@ mod tests { ); let main_c_filepath = std::env::temp_dir().with_file_name(format!( "{}-test-{}.c", - $code_tup.1, + "compiled_function", std::time::SystemTime::now().elapsed().unwrap().as_nanos() )); let binary_filepath = std::env::temp_dir().with_file_name(format!( "{}-test-{}", - $code_tup.1, + "compiled_function", std::time::SystemTime::now().elapsed().unwrap().as_nanos() )); @@ -2078,13 +2078,6 @@ mod tests { "(access-transpose (access-tensor t) (list {}))", permutation.iter().map(|x| x.to_string()).join(" ") ), - ( - { - let hw_map: HashMap = HashMap::default(); - hw_map - }, - "tranpose" - ), input_transposed, "" ); @@ -2125,13 +2118,6 @@ mod tests { "(access-concatenate (access-tensor t0) (access-tensor t1) {})", concat_axis ), - ( - { - let hw_map: HashMap = HashMap::default(); - hw_map - }, - "concatenate" - ), concatted, "" ); @@ -2310,19 +2296,12 @@ int main() {{ "(access-pad (access-tensor t) {} {} {} {})", pad_type, pad_axis, pad_before, pad_after ), - ( - { - let hw_map: HashMap = HashMap::default(); - hw_map - }, - "concatenate" - ), padded, "" ); } #[test] - fn test() { + fn slice() { let input_list = vec![ ( "t", @@ -2361,146 +2340,11 @@ int main() {{ (access-slice (access-tensor t) {} {} {})", slice_axis, low, high ), - ( - { - let hw_map: HashMap = HashMap::default(); - hw_map - }, - "slice" - ), sliced, "" ); } - #[test] - fn slice() { - let shape = vec![32, 7, 100, 3]; - let slice_axis = 2; - let low = 5; - let high = 83; - - let input = ndarray::ArrayD::from_shape_vec( - shape.clone(), - (0..shape.iter().product::()).collect(), - ) - .unwrap(); - - let mut slices = Vec::from_iter( - std::iter::repeat(SliceOrIndex::Slice { - start: 0, - end: None, - step: 1, - }) - .take(shape.len()), - ); - slices[slice_axis] = SliceOrIndex::Slice { - start: low, - end: Some(high), - step: 1, - }; - let sliced = input.slice( - &SliceInfo::, ndarray::IxDyn>::new(slices) - .unwrap() - .as_ref(), - ); - - let expr = RecExpr::from_str( - format!( - " -(access-slice (access-tensor t) {} {} {})", - slice_axis, low, high - ) - .as_str(), - ) - .unwrap(); - - let mut map = HashMap::default(); - map.insert("t".to_string(), shape.clone()); - - let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); - let id = egraph.add_expr(&expr); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "slice", - "", - &vec!["t"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include - -{} -{} -{} -{} - -int main() {{ - slice(out, a); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)a_sliced)[i] == ((float*)out)[i]); - }} -}} -", - c_assignment_string("", "a", DType::Fp32, &input.view()), - c_assignment_string("", "a_sliced", DType::Fp32, &sliced.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(sliced.shape()).view() - ), - code, - sliced.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "slice-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "slice-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - } + #[test] fn access_windows() { From 503c50f51517d2a68a89c35664070cdb9656d7eb Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Tue, 27 Jul 2021 10:57:36 -0700 Subject: [PATCH 05/11] finished first half of codegen test macro --- src/codegen.rs | 319 +++++++++++++++++++++++++++---------------------- 1 file changed, 175 insertions(+), 144 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index 6bd6ed7fba..da3638bd8f 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -1939,7 +1939,7 @@ mod tests { relay_outputs } macro_rules! codegen_test { - ($env: expr, $code_expr: expr, $ground_truth : expr, $c_code : expr) => { + ($env : expr, $code_expr : expr, $hw_map : expr, $ground_truth : expr, $c_code : expr) => { let mut map: HashMap> = HashMap::default(); let env = $env; let mut names = Vec::new(); @@ -1962,7 +1962,7 @@ mod tests { let code = codegen( &egraph, id, - &HashMap::default(), + &$hw_map, "compiled_function", "", &names, @@ -1998,7 +1998,7 @@ mod tests { format!( "{}{}{}", acc, - c_assignment_string("", arg.0, DType::Fp32, &arg.1.view()), + c_assignment_string("", arg.0, DType::Fp32, &arg.1.clone().into_dyn().view()), "\n" ) }), @@ -2059,6 +2059,132 @@ mod tests { .expect("Could not convert stderr to UTF8") ); }; + ($env : expr, $code_expr : expr, $ground_truth : expr, $c_code : expr) => { + let mut map: HashMap> = HashMap::default(); + let env = $env; + let mut names = Vec::new(); + let expr: RecExpr = RecExpr::from_str($code_expr.as_str()).unwrap(); + if (env.len() == 1) { + let name = env[0].0; + map.insert(name.to_string(), env[0].1.shape().to_vec()); + names.push(name); + } else { + for n in 0..(env.len()) { + let name = env[n].0.clone(); + let shape = env[n].1.shape().to_vec().clone(); + + map.insert(name.to_string(), shape); + names.push(name); + } + } + let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); + let id = egraph.add_expr(&expr); + let mut hw_map = HashMap::default(); + hw_map.insert(id, 0); + let code = codegen( + &egraph, + id, + &hw_map, + "compiled_function", + "", + &names, + &generate_worklist_for_codegen(&egraph, id), + true, + ); + let ground_truth = $ground_truth; + let c_func = format!( + " + int main() {{ + {}(out{}); + for (int i = 0; i < {}; i++) {{ + assert(((float*)ground_truth)[i] == ((float*)out)[i]); + }} + }} + ", + "compiled_function", + names + .iter() + .fold(String::new(), |acc, &arg| acc + "," + arg), + ground_truth.shape().iter().product::() + ); + let declarations = format!( + " + #include + #include \"{}\" + {} + {} + {} + {}", + $c_code, + env.iter().fold(String::new(), |acc, arg| { + format!( + "{}{}{}", + acc, + c_assignment_string("", arg.0, DType::Fp32, &arg.1.clone().into_dyn().view()), + "\n" + ) + }), + c_assignment_string("", "ground_truth", DType::Fp32, &ground_truth.view()), + c_assignment_string( + "", + "out", + DType::Fp32, + &ndarray::ArrayD::::zeros(ground_truth.shape()).view() + ), + code + ); + + let main_code = format!( + " + {} + {}", + declarations, c_func + ); + let main_c_filepath = std::env::temp_dir().with_file_name(format!( + "{}-test-{}.c", + "compiled_function", + std::time::SystemTime::now().elapsed().unwrap().as_nanos() + )); + let binary_filepath = std::env::temp_dir().with_file_name(format!( + "{}-test-{}", + "compiled_function", + std::time::SystemTime::now().elapsed().unwrap().as_nanos() + )); + + File::create(&main_c_filepath) + .unwrap() + .write_all(main_code.as_bytes()) + .unwrap(); + + let result = Command::new("gcc") + .arg("-Werror") + .arg("-g") + .arg("-o") + .arg(&binary_filepath) + .arg(&main_c_filepath) + .output() + .unwrap(); + + assert!( + result.status.success(), + "{}", + std::str::from_utf8(result.stderr.as_slice()) + .expect("Could not convert stderr to UTF8") + ); + + let result = Command::new(&binary_filepath).output().unwrap(); + + assert!( + result.status.success(), + "{}", + std::str::from_utf8(result.stderr.as_slice()) + .expect("Could not convert stderr to UTF8") + ); + }; + ($env: expr, $code_expr: expr, $ground_truth : expr) => { + codegen_test!($env, $code_expr, HashMap::default(), $ground_truth, ""); + }; + } #[test] fn tranpose() { @@ -2078,8 +2204,7 @@ mod tests { "(access-transpose (access-tensor t) (list {}))", permutation.iter().map(|x| x.to_string()).join(" ") ), - input_transposed, - "" + input_transposed ); } @@ -2118,147 +2243,55 @@ mod tests { "(access-concatenate (access-tensor t0) (access-tensor t1) {})", concat_axis ), - concatted, - "" + concatted ); } - #[test] fn systolic_array() { - let shape0 = vec![2, 10]; - let shape1 = vec![10, 15]; - - let input0 = ndarray::ArrayD::from_shape_vec( - shape0.clone(), - (0..shape0.iter().product::()).collect(), - ) - .unwrap() - .into_dimensionality::() - .unwrap(); - let input1 = ndarray::ArrayD::from_shape_vec( - shape1.clone(), - (0..shape1.iter().product::()).collect(), - ) - .unwrap() - .into_dimensionality::() - .unwrap(); - let multiplied = input0.dot(&input1).into_dyn(); - - let expr = RecExpr::from_str( - " -(systolic-array 10 15 - (access (access-tensor t0) 1) - (access (access-tensor t1) 0) -)", - ) - .unwrap(); - - let mut map = HashMap::default(); - map.insert("t0".to_string(), shape0.clone()); - map.insert("t1".to_string(), shape1.clone()); - - let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); - let id = egraph.add_expr(&expr); - - let mut hw_map = HashMap::default(); - hw_map.insert(id, 0); - - let code = codegen( - &egraph, - id, - &hw_map, - "systolic_array", - "", - &vec!["t0", "t1"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include -#include \"{}\" - -{} -{} -{} -{} -{} - -int main() {{ - systolic_array(out, t0, t1); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)result)[i] == ((float*)out)[i]); - }} -}} -", - PathBuf::from_str( - format!( - "{}/{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "data", - "codegen-mlp", - "rtml_systolic_array_weight_stationary.c" + let input_list = vec![ + ( + "t0", + ndarray::ArrayD::from_shape_vec( + vec![2, 10].clone(), + (0..vec![2, 10].iter().product::()).collect(), ) - .as_str() - ) - .unwrap() - .to_string_lossy(), - c_assignment_string("", "t0", DType::Fp32, &input0.into_dyn().view()), - c_assignment_string("", "t1", DType::Fp32, &input1.into_dyn().view()), - c_assignment_string("", "result", DType::Fp32, &multiplied.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(multiplied.shape()).view() + .unwrap() + .into_dimensionality::() + .unwrap() ), - code, - multiplied.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "systolic-array-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "systolic-array-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); + ( + "t1", + ndarray::ArrayD::from_shape_vec( + vec![10, 15].clone(), + (0..vec![10, 15].iter().product::()).collect(), + ) + .unwrap() + .into_dimensionality::() + .unwrap() + )]; + let multiplied = input_list[0].1.clone().dot(&input_list[1].1.clone()).into_dyn(); + codegen_test!( + input_list.clone(), + format!(" + (systolic-array 10 15 + (access (access-tensor t0) 1) + (access (access-tensor t1) 0) + )" + ), + multiplied, + PathBuf::from_str( + format!( + "{}/{}/{}/{}", + env!("CARGO_MANIFEST_DIR"), + "data", + "codegen-mlp", + "rtml_systolic_array_weight_stationary.c" + ) + .as_str() + ) + .unwrap() + .to_string_lossy() + ); } #[test] fn pad() { @@ -2296,8 +2329,7 @@ int main() {{ "(access-pad (access-tensor t) {} {} {} {})", pad_type, pad_axis, pad_before, pad_after ), - padded, - "" + padded ); } #[test] @@ -2340,8 +2372,7 @@ int main() {{ (access-slice (access-tensor t) {} {} {})", slice_axis, low, high ), - sliced, - "" + sliced ); } From b059efd3615fd6e156e8e34e63dd0f8b1b96f9ce Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Tue, 27 Jul 2021 14:57:41 -0700 Subject: [PATCH 06/11] fmt --- src/codegen.rs | 82 +++++++++++++++++++++++++++++--------------------- 1 file changed, 47 insertions(+), 35 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index da3638bd8f..9cbf03907f 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -1998,7 +1998,12 @@ mod tests { format!( "{}{}{}", acc, - c_assignment_string("", arg.0, DType::Fp32, &arg.1.clone().into_dyn().view()), + c_assignment_string( + "", + arg.0, + DType::Fp32, + &arg.1.clone().into_dyn().view() + ), "\n" ) }), @@ -2120,7 +2125,12 @@ mod tests { format!( "{}{}{}", acc, - c_assignment_string("", arg.0, DType::Fp32, &arg.1.clone().into_dyn().view()), + c_assignment_string( + "", + arg.0, + DType::Fp32, + &arg.1.clone().into_dyn().view() + ), "\n" ) }), @@ -2184,7 +2194,6 @@ mod tests { ($env: expr, $code_expr: expr, $ground_truth : expr) => { codegen_test!($env, $code_expr, HashMap::default(), $ground_truth, ""); }; - } #[test] fn tranpose() { @@ -2250,14 +2259,14 @@ mod tests { fn systolic_array() { let input_list = vec![ ( - "t0", + "t0", ndarray::ArrayD::from_shape_vec( vec![2, 10].clone(), (0..vec![2, 10].iter().product::()).collect(), ) .unwrap() .into_dimensionality::() - .unwrap() + .unwrap(), ), ( "t1", @@ -2267,31 +2276,37 @@ mod tests { ) .unwrap() .into_dimensionality::() - .unwrap() - )]; - let multiplied = input_list[0].1.clone().dot(&input_list[1].1.clone()).into_dyn(); - codegen_test!( - input_list.clone(), - format!(" + .unwrap(), + ), + ]; + let multiplied = input_list[0] + .1 + .clone() + .dot(&input_list[1].1.clone()) + .into_dyn(); + codegen_test!( + input_list.clone(), + format!( + " (systolic-array 10 15 (access (access-tensor t0) 1) (access (access-tensor t1) 0) )" - ), - multiplied, - PathBuf::from_str( - format!( - "{}/{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "data", - "codegen-mlp", - "rtml_systolic_array_weight_stationary.c" - ) - .as_str() + ), + multiplied, + PathBuf::from_str( + format!( + "{}/{}/{}/{}", + env!("CARGO_MANIFEST_DIR"), + "data", + "codegen-mlp", + "rtml_systolic_array_weight_stationary.c" ) - .unwrap() - .to_string_lossy() - ); + .as_str() + ) + .unwrap() + .to_string_lossy() + ); } #[test] fn pad() { @@ -2334,16 +2349,14 @@ mod tests { } #[test] fn slice() { - let input_list = vec![ - ( - "t", - ndarray::ArrayD::from_shape_vec( - vec![32, 7, 100, 3].clone(), - (0..vec![32, 7, 100, 3].iter().product::()).collect(), - ) - .unwrap() + let input_list = vec![( + "t", + ndarray::ArrayD::from_shape_vec( + vec![32, 7, 100, 3].clone(), + (0..vec![32, 7, 100, 3].iter().product::()).collect(), ) - ]; + .unwrap(), + )]; let slice_axis = 2; let low = 5; let high = 83; @@ -2375,7 +2388,6 @@ mod tests { sliced ); } - #[test] fn access_windows() { From 59c12142e499f0e7a2d96c6a36a65ed2af91040f Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Thu, 26 Aug 2021 17:14:48 -0700 Subject: [PATCH 07/11] Added more tests; fixed merge issue? --- src/codegen.rs | 255 ++++++++++--------------------------------------- 1 file changed, 53 insertions(+), 202 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index 9cbf03907f..7af9da937f 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -2388,19 +2388,20 @@ mod tests { sliced ); } - #[test] fn access_windows() { let shape = vec![3, 50, 27, 4]; let filters_shape = vec![2, 9, 22, 3]; let stride = vec![1, 10, 3, 1]; - let input = ndarray::ArrayD::from_shape_vec( - shape.clone(), - (0i64..shape.iter().map(|v| *v as i64).product::()).collect(), - ) - .unwrap(); - + let input_list = vec![( + "t", + ndarray::ArrayD::from_shape_vec( + shape.clone(), + (0i64..shape.iter().map(|v| *v as i64).product::()).collect(), + ) + .unwrap(), + )]; let expr = RecExpr::from_str( format!( " @@ -2424,130 +2425,51 @@ mod tests { .as_str(), ) .unwrap(); - - let mut map = HashMap::default(); - map.insert("t".to_string(), shape.clone()); - let mut env = HashMap::default(); - env.insert("t", input.clone()); + env.insert("t", input_list[0].1.clone()); let out = - match crate::language::interpreter::interpret(&expr, expr.as_ref().len() - 1, &env) { - crate::language::interpreter::Value::Access(a) => a, - _ => panic!(), - }; - - let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); - let id = egraph.add_expr(&expr); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "access_windows", - "", - &vec!["t"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include - -{} -{} -{} -{} - -int main() {{ - access_windows(out, a); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)a_windows)[i] == ((float*)out)[i]); - }} -}} -", - c_assignment_string("", "a", DType::Fp32, &input.view()), - c_assignment_string("", "a_windows", DType::Fp32, &out.tensor.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(out.tensor.shape()).view() + match crate::language::interpreter::interpret(&expr, expr.as_ref().len() - 1, &env) { + crate::language::interpreter::Value::Access(a) => a, + _ => panic!(), + }; + codegen_test!( + input_list.clone(), + format!( + " +(access-windows + (access (access-tensor t) {access_axis}) + (shape {filter_shapes}) + (shape {strides}) +)", + access_axis = shape.len(), + filter_shapes = filters_shape + .iter() + .map(|i| i.to_string()) + .collect::>() + .join(" "), + strides = stride + .iter() + .map(|i| i.to_string()) + .collect::>() + .join(" "), ), - code, - out.tensor.shape().iter().product::() + out.tensor ); - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "access-windows-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "access-windows-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); } - #[test] fn access_flatten() { let shape = vec![3, 50, 27, 4]; let access_axis = 2; - - let input = ndarray::ArrayD::from_shape_vec( - shape.clone(), - (0..shape.iter().product::()).collect(), - ) - .unwrap(); - - let expr = RecExpr::from_str( - format!( - " -(access-flatten - (access (access-tensor t) {access_axis}) -)", - access_axis = access_axis + let input_list = vec![( + "t", + ndarray::ArrayD::from_shape_vec( + shape.clone(), + (0..shape.iter().product::()).collect(), ) - .as_str(), - ) - .unwrap(); - - let mut map = HashMap::default(); - map.insert("t".to_string(), shape.clone()); - - let out = input + .unwrap() + )]; + let out = input_list[0].1 .view() .into_shape(( shape[0..access_axis].iter().product::(), @@ -2555,92 +2477,21 @@ int main() {{ )) .unwrap() .into_dyn(); - - let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); - let id = egraph.add_expr(&expr); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "access_flatten", - "", - &vec!["t"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include - -{} -{} -{} -{} - -int main() {{ - access_flatten(out, a); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)a_flattened)[i] == ((float*)out)[i]); - }} -}} -", - c_assignment_string("", "a", DType::Fp32, &input.view()), - c_assignment_string("", "a_flattened", DType::Fp32, &out.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(out.shape()).view() + codegen_test!( + input_list.clone(), + format!( + " +(access-flatten + (access (access-tensor t) {access_axis}) +)", + access_axis = access_axis ), - code, - out.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "access-flatten-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "access-flatten-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") + out ); } + + #[test] #[should_panic] fn extract_monolithic_panic() { From 585286f51849cc3404de8ad64a182aa0925b3ad6 Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Thu, 26 Aug 2021 17:16:06 -0700 Subject: [PATCH 08/11] fmt --- src/codegen.rs | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index 7af9da937f..b9e2d80753 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -2428,10 +2428,10 @@ mod tests { let mut env = HashMap::default(); env.insert("t", input_list[0].1.clone()); let out = - match crate::language::interpreter::interpret(&expr, expr.as_ref().len() - 1, &env) { - crate::language::interpreter::Value::Access(a) => a, - _ => panic!(), - }; + match crate::language::interpreter::interpret(&expr, expr.as_ref().len() - 1, &env) { + crate::language::interpreter::Value::Access(a) => a, + _ => panic!(), + }; codegen_test!( input_list.clone(), format!( @@ -2455,7 +2455,6 @@ mod tests { ), out.tensor ); - } #[test] fn access_flatten() { @@ -2467,9 +2466,10 @@ mod tests { shape.clone(), (0..shape.iter().product::()).collect(), ) - .unwrap() + .unwrap(), )]; - let out = input_list[0].1 + let out = input_list[0] + .1 .view() .into_shape(( shape[0..access_axis].iter().product::(), @@ -2490,8 +2490,6 @@ mod tests { ); } - - #[test] #[should_panic] fn extract_monolithic_panic() { From 145f1cab0f0ad7b22bf37cc6ef2175f29c6dd838 Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Tue, 7 Sep 2021 10:28:53 -0700 Subject: [PATCH 09/11] finished first half of tests (all non-relay, non-panic tests) --- src/codegen.rs | 161 ++++++++++++------------------------------------- 1 file changed, 38 insertions(+), 123 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index b9e2d80753..a5cb961948 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -2525,77 +2525,45 @@ mod tests { let (_hw_map, _hw_design) = create_hardware_design_monolithic(&egraph, (32, 32)); } - #[test] fn systolic_array_with_blocking() { - let shape0 = vec![2, 10]; - let shape1 = vec![10, 15]; - - let input0 = ndarray::ArrayD::from_shape_vec( - shape0.clone(), - (0..shape0.iter().product::()).collect(), - ) - .unwrap() - .into_dimensionality::() - .unwrap(); - let input1 = ndarray::ArrayD::from_shape_vec( - shape1.clone(), - (0..shape1.iter().product::()).collect(), - ) - .unwrap() - .into_dimensionality::() - .unwrap(); - let multiplied = input0.dot(&input1).into_dyn(); - - let expr = RecExpr::from_str( - " -(systolic-array-with-blocking 2 5 - (access (access-tensor t0) 1) - (access (access-tensor t1) 0) -)", - ) - .unwrap(); - - let mut map = HashMap::default(); - map.insert("t0".to_string(), shape0.clone()); - map.insert("t1".to_string(), shape1.clone()); - - let mut egraph = EGraph::new(MyAnalysis { name_to_shape: map }); - let id = egraph.add_expr(&expr); - - let mut hw_map = HashMap::default(); - hw_map.insert(id, 0); - - let code = codegen( - &egraph, - id, - &hw_map, - "systolic_array_with_blocking", - "", - &vec!["t0", "t1"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include -#include \"{}\" - -{} -{} -{} -{} -{} - -int main() {{ - systolic_array_with_blocking(out, t0, t1); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)result)[i] == ((float*)out)[i]); - }} -}} -", + let input_list = vec![ + ( + "t0", + ndarray::ArrayD::from_shape_vec( + vec![2, 10].clone(), + (0..vec![2, 10].iter().product::()).collect(), + ) + .unwrap() + .into_dimensionality::() + .unwrap(), + ), + ( + "t1", + ndarray::ArrayD::from_shape_vec( + vec![10, 15].clone(), + (0..vec![10, 15].iter().product::()).collect(), + ) + .unwrap() + .into_dimensionality::() + .unwrap(), + ), + ]; + let multiplied = input_list[0] + .1 + .clone() + .dot(&input_list[1].1.clone()) + .into_dyn(); + codegen_test!( + input_list.clone(), + format!( + " + (systolic-array-with-blocking 2 5 + (access (access-tensor t0) 1) + (access (access-tensor t1) 0) + )" + ), + multiplied, PathBuf::from_str( format!( "{}/{}/{}/{}", @@ -2607,60 +2575,7 @@ int main() {{ .as_str() ) .unwrap() - .to_string_lossy(), - c_assignment_string("", "t0", DType::Fp32, &input0.into_dyn().view()), - c_assignment_string("", "t1", DType::Fp32, &input1.into_dyn().view()), - c_assignment_string("", "result", DType::Fp32, &multiplied.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(multiplied.shape()).view() - ), - code, - multiplied.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "systolic-array-with-blocking-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "systolic-array-with-blocking-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") + .to_string_lossy() ); } From 689e9bf00e7f3db652dfde2f1b97752c66c88742 Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Wed, 19 Jan 2022 10:41:12 -0800 Subject: [PATCH 10/11] finished first part of relay_test macro --- src/codegen.rs | 280 +++++++++++++++++++++++++------------------------ 1 file changed, 143 insertions(+), 137 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index a5cb961948..829b3bce14 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -1938,6 +1938,135 @@ mod tests { } relay_outputs } + macro_rules! relay_op_test { + ($name : expr, $relay_code : expr, $relay_op : expr, 2, $parameters : expr) => { + let relay = $relay_code; + let module = tvm::ir::module::IRModule::parse("", relay).unwrap(); + let (expr, shapes_vec) = crate::language::from_relay::from_relay( + &module, + true, + &vec![$relay_op], + ); + let mut env: HashMap> = HashMap::default(); + for (k, v) in &shapes_vec { + env.insert(k.clone(), v.clone()); + } + let mut egraph = EGraph::new(MyAnalysis { + name_to_shape: env.clone(), + }); + + let id = egraph.add_expr(&expr); + + let x_input = ndarray::ArrayD::from_shape_vec( + env.get("x").unwrap().clone(), + (0..env.get("x").unwrap().iter().product::()).collect(), + ) + .unwrap(); + let y_input = ndarray::ArrayD::from_shape_vec( + env.get("y").unwrap().clone(), + (0..env.get("y").unwrap().iter().product::()).collect(), + ) + .unwrap(); + let result = (&x_input + &y_input).into_dyn(); + let code = codegen( + &egraph, + id, + &HashMap::default(), + "relay_biasadd", + "", + &$parameters, + &generate_worklist_for_codegen(&egraph, id), + true, + ); + let main_code = format!( + " + #include + #include \"{}\" + + {} + {} + {} + {} + {} + + int main() {{ + relay_biasadd(out, x, y); + + for (int i = 0; i < {}; i++) {{ + assert(((float*)result)[i] == ((float*)out)[i]); + }} + }} + ", + PathBuf::from_str( + format!( + "{}/{}/{}", + env!("CARGO_MANIFEST_DIR"), + "c-files", + "relay-op-implementations.c" + ) + .as_str() + ) + .unwrap() + .to_string_lossy(), + c_assignment_string("", "x", DType::Fp32, &x_input.into_dyn().view()), + c_assignment_string("", "y", DType::Fp32, &y_input.into_dyn().view()), + c_assignment_string("", "result", DType::Fp32, &result.view()), + c_assignment_string( + "", + "out", + DType::Fp32, + &ndarray::ArrayD::::zeros(result.shape()).view() + ), + code, + result.shape().iter().product::() + ); + + let main_c_filepath = std::env::temp_dir().with_file_name(format!( + "relay-op-{}-test-{}.c", + $name, + std::time::SystemTime::now().elapsed().unwrap().as_nanos() + )); + println!("{}", main_c_filepath.to_string_lossy()); + + let binary_filepath = std::env::temp_dir().with_file_name(format!( + "relay-op-{}-test-{}", + $name, + std::time::SystemTime::now().elapsed().unwrap().as_nanos() + )); + println!("{}", binary_filepath.to_string_lossy()); + + File::create(&main_c_filepath) + .unwrap() + .write_all(main_code.as_bytes()) + .unwrap(); + + let result = Command::new("gcc") + .arg("-Werror") + .arg("-g") + .arg("-o") + .arg(&binary_filepath) + .arg(&main_c_filepath) + .arg("-lm") + .output() + .unwrap(); + + assert!( + result.status.success(), + "{}", + std::str::from_utf8(result.stderr.as_slice()) + .expect("Could not convert stderr to UTF8") + ); + + let result = Command::new(&binary_filepath).output().unwrap(); + + assert!( + result.status.success(), + "{}", + std::str::from_utf8(result.stderr.as_slice()) + .expect("Could not convert stderr to UTF8") + ); + } + } macro_rules! codegen_test { ($env : expr, $code_expr : expr, $hw_map : expr, $ground_truth : expr, $c_code : expr) => { let mut map: HashMap> = HashMap::default(); @@ -2195,6 +2324,8 @@ mod tests { codegen_test!($env, $code_expr, HashMap::default(), $ground_truth, ""); }; } + + #[test] fn tranpose() { let permutation = vec![3, 1, 0, 2]; @@ -2328,7 +2459,6 @@ mod tests { let mut pad_after_shape = input_list[0].1.shape().to_vec().clone(); pad_after_shape[pad_axis] = pad_after; assert!(pad_type == PadType::ZeroPadding); - let padded = ndarray::stack( ndarray::Axis(pad_axis), &[ @@ -2336,8 +2466,7 @@ mod tests { input_list[0].1.view(), ndarray::ArrayD::zeros(pad_after_shape).view(), ], - ) - .unwrap(); + ).unwrap(); codegen_test!( input_list.clone(), format!( @@ -2716,142 +2845,19 @@ int main() {{ .expect("Could not convert stderr to UTF8") ); } - #[test] fn relay_op_biasadd() { - let relay = r#" -#[version = "0.0.5"] -def @main(%x: Tensor[(1, 1000), float32], %y: Tensor[(1000), float32]) { - nn.bias_add(%x, %y) -} -"#; - - let module = tvm::ir::module::IRModule::parse("", relay).unwrap(); - - let (expr, shapes_vec) = crate::language::from_relay::from_relay( - &module, - true, - &vec![crate::language::RelayOperator::RelayBiasAdd], - ); - - let mut env = HashMap::default(); - for (k, v) in &shapes_vec { - env.insert(k.clone(), v.clone()); - } - - let mut egraph = EGraph::new(MyAnalysis { - name_to_shape: env.clone(), - }); - - let id = egraph.add_expr(&expr); - - let x_input = ndarray::ArrayD::from_shape_vec( - env.get("x").unwrap().clone(), - (0..env.get("x").unwrap().iter().product::()).collect(), - ) - .unwrap(); - let y_input = ndarray::ArrayD::from_shape_vec( - env.get("y").unwrap().clone(), - (0..env.get("y").unwrap().iter().product::()).collect(), - ) - .unwrap(); - let result = (&x_input + &y_input).into_dyn(); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "relay_biasadd", - "", - &vec!["x", "y"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include -#include \"{}\" - -{} -{} -{} -{} -{} - -int main() {{ - relay_biasadd(out, x, y); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)result)[i] == ((float*)out)[i]); - }} -}} -", - PathBuf::from_str( - format!( - "{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "c-files", - "relay-op-implementations.c" - ) - .as_str() - ) - .unwrap() - .to_string_lossy(), - c_assignment_string("", "x", DType::Fp32, &x_input.into_dyn().view()), - c_assignment_string("", "y", DType::Fp32, &y_input.into_dyn().view()), - c_assignment_string("", "result", DType::Fp32, &result.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(result.shape()).view() - ), - code, - result.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-biasadd-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-biasadd-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .arg("-lm") - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") + relay_op_test!( + "relay_op_biasadd", + r#" + #[version = "0.0.5"] + def @main(%x: Tensor[(1, 1000), float32], %y: Tensor[(1000), float32]) { + nn.bias_add(%x, %y) + } + "#, + crate::language::RelayOperator::RelayBiasAdd, + 2, + vec!["x", "y"] ); } From f8ab1408069c78c72be7136099a1c1e4191b207d Mon Sep 17 00:00:00 2001 From: Vishal Canumalla Date: Mon, 24 Jan 2022 21:53:08 -0800 Subject: [PATCH 11/11] finished first pass of codegen test macro(s) --- src/codegen.rs | 1251 ++++++++++-------------------------------------- 1 file changed, 250 insertions(+), 1001 deletions(-) diff --git a/src/codegen.rs b/src/codegen.rs index 829b3bce14..f8888315be 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -1785,7 +1785,11 @@ mod tests { use std::path::PathBuf; use std::process::Command; use std::str::FromStr; - + // for macro call + use ndarray::ArrayBase; + use ndarray::Dim; + use ndarray::IxDynImpl; + use ndarray::OwnedRepr; fn run_relay( env: &HashMap>, shapes_vec: &Vec<(String, Vec)>, @@ -1939,14 +1943,77 @@ mod tests { relay_outputs } macro_rules! relay_op_test { - ($name : expr, $relay_code : expr, $relay_op : expr, 2, $parameters : expr) => { + // Macro for calling relay_op tests. There are several pattern matches here. In order, here is how to use them + // The first pattern match is for tests that utilize a `value_env`. The user inputs the required fields. For tests + // where the relay code used to form the module and the result are the same, use the same relay code expression for + // `relay_code` and `run_relay` + ($name: expr, $relay_code : expr, $run_relay : expr, $relay_op : expr, $seed : expr, $lower_bound : expr, $upper_bound : expr, $parameters : expr) => { let relay = $relay_code; + const SEED: u64 = $seed; + let mut tensor_rng = SmallRng::seed_from_u64(SEED); let module = tvm::ir::module::IRModule::parse("", relay).unwrap(); - let (expr, shapes_vec) = crate::language::from_relay::from_relay( - &module, + let (expr, shapes_vec) = + crate::language::from_relay::from_relay(&module, true, &vec![$relay_op]); + let mut env: HashMap> = HashMap::default(); + let mut value_env: HashMap, Dim>> = + HashMap::default(); + for (k, v) in &shapes_vec { + env.insert(k.clone(), v.clone()); + value_env.insert( + k.clone(), + ndarray::ArrayD::::random_using( + v.clone(), + Uniform::new($lower_bound, $upper_bound), // prevent NaN results + &mut tensor_rng, + ), + ); + } + let mut egraph = EGraph::new(MyAnalysis { + name_to_shape: env.clone(), + }); + + let id = egraph.add_expr(&expr); + let result = run_relay(&value_env, &shapes_vec, $run_relay); + let parameters = $parameters; + let code = codegen( + &egraph, + id, + &HashMap::default(), + "compiled_function", + "", + ¶meters, + &generate_worklist_for_codegen(&egraph, id), true, - &vec![$relay_op], ); + relay_op_test!($name, value_env, parameters, code, result); + }; + ($name : expr, $value_env : expr, $parameters : expr, $c_code : expr, $result : expr) => { + let mut var_assignment_str = "".to_string(); + let mut var = "".to_string(); + for i in 0..($parameters.len()) { + let cur = c_assignment_string( + "", + $parameters[i], + DType::Fp32, + &$value_env.get($parameters[i]).unwrap().view(), + ); + var_assignment_str = format!("{}{}", var_assignment_str, cur); + var = format!("{}{}{}", var, ", ", $parameters[i]); + } + relay_op_test!( + $name, + var_assignment_str, + var, + "fabs(((float*)result)[i] - ((float*)out)[i]) < 0.00001", + $c_code, + $result + ); + }; + ($name : expr, $relay_code : expr, $relay_op : expr) => { + let relay = $relay_code; + let module = tvm::ir::module::IRModule::parse("", relay).unwrap(); + let (expr, shapes_vec) = + crate::language::from_relay::from_relay(&module, true, &vec![$relay_op]); let mut env: HashMap> = HashMap::default(); for (k, v) in &shapes_vec { env.insert(k.clone(), v.clone()); @@ -1954,7 +2021,7 @@ mod tests { let mut egraph = EGraph::new(MyAnalysis { name_to_shape: env.clone(), }); - + let id = egraph.add_expr(&expr); let x_input = ndarray::ArrayD::from_shape_vec( @@ -1972,12 +2039,27 @@ mod tests { &egraph, id, &HashMap::default(), - "relay_biasadd", + "compiled_function", "", - &$parameters, + &vec!["x", "y"], &generate_worklist_for_codegen(&egraph, id), true, ); + let var_assignment_str = format!( + "{}{}", + c_assignment_string("", "x", DType::Fp32, &x_input.into_dyn().view()), + c_assignment_string("", "y", DType::Fp32, &y_input.into_dyn().view()) + ); + relay_op_test!( + $name, + var_assignment_str, + ", x, y", + "((float*)result)[i] == ((float*)out)[i]", + code, + result + ); + }; + ($name : expr, $var_assignment_str: expr, $var : expr, $c_string: expr, $c_code : expr, $result : expr) => { let main_code = format!( " #include @@ -1987,13 +2069,12 @@ mod tests { {} {} {} - {} int main() {{ - relay_biasadd(out, x, y); + compiled_function(out {}); for (int i = 0; i < {}; i++) {{ - assert(((float*)result)[i] == ((float*)out)[i]); + assert({}); }} }} ", @@ -2008,38 +2089,39 @@ mod tests { ) .unwrap() .to_string_lossy(), - c_assignment_string("", "x", DType::Fp32, &x_input.into_dyn().view()), - c_assignment_string("", "y", DType::Fp32, &y_input.into_dyn().view()), - c_assignment_string("", "result", DType::Fp32, &result.view()), + $var_assignment_str, + c_assignment_string("", "result", DType::Fp32, &$result.view()), c_assignment_string( "", "out", DType::Fp32, - &ndarray::ArrayD::::zeros(result.shape()).view() + &ndarray::ArrayD::::zeros($result.shape()).view() ), - code, - result.shape().iter().product::() + $c_code, + $var, + $result.shape().iter().product::(), + $c_string ); - + let main_c_filepath = std::env::temp_dir().with_file_name(format!( "relay-op-{}-test-{}.c", - $name, + $name, std::time::SystemTime::now().elapsed().unwrap().as_nanos() )); println!("{}", main_c_filepath.to_string_lossy()); - + let binary_filepath = std::env::temp_dir().with_file_name(format!( "relay-op-{}-test-{}", $name, std::time::SystemTime::now().elapsed().unwrap().as_nanos() )); println!("{}", binary_filepath.to_string_lossy()); - + File::create(&main_c_filepath) .unwrap() .write_all(main_code.as_bytes()) .unwrap(); - + let result = Command::new("gcc") .arg("-Werror") .arg("-g") @@ -2049,23 +2131,23 @@ mod tests { .arg("-lm") .output() .unwrap(); - + assert!( result.status.success(), "{}", std::str::from_utf8(result.stderr.as_slice()) .expect("Could not convert stderr to UTF8") ); - + let result = Command::new(&binary_filepath).output().unwrap(); - + assert!( result.status.success(), "{}", std::str::from_utf8(result.stderr.as_slice()) .expect("Could not convert stderr to UTF8") ); - } + }; } macro_rules! codegen_test { ($env : expr, $code_expr : expr, $hw_map : expr, $ground_truth : expr, $c_code : expr) => { @@ -2324,7 +2406,6 @@ mod tests { codegen_test!($env, $code_expr, HashMap::default(), $ground_truth, ""); }; } - #[test] fn tranpose() { @@ -2466,7 +2547,8 @@ mod tests { input_list[0].1.view(), ndarray::ArrayD::zeros(pad_after_shape).view(), ], - ).unwrap(); + ) + .unwrap(); codegen_test!( input_list.clone(), format!( @@ -2707,142 +2789,17 @@ mod tests { .to_string_lossy() ); } - #[test] fn relay_op_add() { - let relay = r#" -#[version = "0.0.5"] -def @main(%x: Tensor[(1, 16, 16, 3), float32], %y: Tensor[(1, 1, 3), float32]) { - add(%x, %y) -} -"#; - - let module = tvm::ir::module::IRModule::parse("", relay).unwrap(); - - let (expr, shapes_vec) = crate::language::from_relay::from_relay( - &module, - true, - &vec![crate::language::RelayOperator::RelayAdd], - ); - - let mut env = HashMap::default(); - for (k, v) in &shapes_vec { - env.insert(k.clone(), v.clone()); - } - - let mut egraph = EGraph::new(MyAnalysis { - name_to_shape: env.clone(), - }); - - let id = egraph.add_expr(&expr); - - let x_input = ndarray::ArrayD::from_shape_vec( - env.get("x").unwrap().clone(), - (0..env.get("x").unwrap().iter().product::()).collect(), - ) - .unwrap(); - let y_input = ndarray::ArrayD::from_shape_vec( - env.get("y").unwrap().clone(), - (0..env.get("y").unwrap().iter().product::()).collect(), - ) - .unwrap(); - let result = (&x_input + &y_input).into_dyn(); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "relay_add", - "", - &vec!["x", "y"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include -#include \"{}\" - -{} -{} -{} -{} -{} - -int main() {{ - relay_add(out, x, y); - - for (int i = 0; i < {}; i++) {{ - assert(((float*)result)[i] == ((float*)out)[i]); - }} -}} -", - PathBuf::from_str( - format!( - "{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "c-files", - "relay-op-implementations.c" - ) - .as_str() - ) - .unwrap() - .to_string_lossy(), - c_assignment_string("", "x", DType::Fp32, &x_input.into_dyn().view()), - c_assignment_string("", "y", DType::Fp32, &y_input.into_dyn().view()), - c_assignment_string("", "result", DType::Fp32, &result.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(result.shape()).view() - ), - code, - result.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-add-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-add-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .arg("-lm") - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") + relay_op_test!( + "relay_op_add", + r#" + #[version = "0.0.5"] + def @main(%x: Tensor[(1, 16, 16, 3), float32], %y: Tensor[(1, 1, 3), float32]) { + add(%x, %y) + } + "#, + crate::language::RelayOperator::RelayAdd ); } #[test] @@ -2855,44 +2812,55 @@ int main() {{ nn.bias_add(%x, %y) } "#, - crate::language::RelayOperator::RelayBiasAdd, - 2, - vec!["x", "y"] + crate::language::RelayOperator::RelayBiasAdd ); } - #[test] fn relay_op_batchnorm() { + relay_op_test!( + "relay_op_batchnorm", + r#" + #[version = "0.0.5"] + def @main(%data: Tensor[(1, 2, 2, 16), float32], %bn_gamma: Tensor[(16), float32], %bn_beta: Tensor[(16), float32], %bn_mean: Tensor[(16), float32], %bn_var: Tensor[(16), float32]) -> (Tensor[(1, 2, 2, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) { + nn.batch_norm(%data, %bn_gamma, %bn_beta, %bn_mean, %bn_var, axis=3) /* ty=(Tensor[(1, 2, 2, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) */ + } + "#, + r#" + #[version = "0.0.5"] + def @main(%data: Tensor[(1, 2, 2, 16), float32], %bn_gamma: Tensor[(16), float32], %bn_beta: Tensor[(16), float32], %bn_mean: Tensor[(16), float32], %bn_var: Tensor[(16), float32]) -> Tensor[(1, 2, 2, 16), float32] { + %0 = add(%bn_var, 1e-05f /* ty=float32 */) /* ty=Tensor[(16), float32] */; + %1 = sqrt(%0) /* ty=Tensor[(16), float32] */; + %2 = divide(1f /* ty=float32 */, %1) /* ty=Tensor[(16), float32] */; + %3 = multiply(%2, %bn_gamma) /* ty=Tensor[(16), float32] */; + %4 = multiply(%data, %3) /* ty=Tensor[(1, 2, 2, 16), float32] */; + %5 = negative(%bn_mean) /* ty=Tensor[(16), float32] */; + %6 = multiply(%5, %3) /* ty=Tensor[(16), float32] */; + %7 = add(%6, %bn_beta) /* ty=Tensor[(16), float32] */; + add(%4, %7) /* ty=Tensor[(1, 2, 2, 16), float32] */ + } + "#, + crate::language::RelayOperator::RelayBatchNormInference, + 23, + 1f32, + 2f32, + vec!["data", "bn_gamma", "bn_beta", "bn_mean", "bn_var"] + ); + } + #[test] + fn relay_op_softmax() { let relay = r#" -#[version = "0.0.5"] -def @main(%data: Tensor[(1, 2, 2, 16), float32], %bn_gamma: Tensor[(16), float32], %bn_beta: Tensor[(16), float32], %bn_mean: Tensor[(16), float32], %bn_var: Tensor[(16), float32]) -> (Tensor[(1, 2, 2, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) { - nn.batch_norm(%data, %bn_gamma, %bn_beta, %bn_mean, %bn_var, axis=3) /* ty=(Tensor[(1, 2, 2, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) */ -} -"#; - let relay_to_run = r#" -#[version = "0.0.5"] -def @main(%data: Tensor[(1, 2, 2, 16), float32], %bn_gamma: Tensor[(16), float32], %bn_beta: Tensor[(16), float32], %bn_mean: Tensor[(16), float32], %bn_var: Tensor[(16), float32]) -> Tensor[(1, 2, 2, 16), float32] { - %0 = add(%bn_var, 1e-05f /* ty=float32 */) /* ty=Tensor[(16), float32] */; - %1 = sqrt(%0) /* ty=Tensor[(16), float32] */; - %2 = divide(1f /* ty=float32 */, %1) /* ty=Tensor[(16), float32] */; - %3 = multiply(%2, %bn_gamma) /* ty=Tensor[(16), float32] */; - %4 = multiply(%data, %3) /* ty=Tensor[(1, 2, 2, 16), float32] */; - %5 = negative(%bn_mean) /* ty=Tensor[(16), float32] */; - %6 = multiply(%5, %3) /* ty=Tensor[(16), float32] */; - %7 = add(%6, %bn_beta) /* ty=Tensor[(16), float32] */; - add(%4, %7) /* ty=Tensor[(1, 2, 2, 16), float32] */ -} -"#; - // Random number generator for generating random tensors. - const SEED: u64 = 23; - let mut tensor_rng = SmallRng::seed_from_u64(SEED); + #[version = "0.0.5"] + def @main(%data: Tensor[(1,100), float32]) -> Tensor[(1,100), float32] { + nn.softmax(%data) /* ty=Tensor[(1,10), float32] */ + } + "#; let module = tvm::ir::module::IRModule::parse("", relay).unwrap(); let (expr, shapes_vec) = crate::language::from_relay::from_relay( &module, true, - &vec![crate::language::RelayOperator::RelayBatchNormInference], + &vec![crate::language::RelayOperator::RelaySoftmax], ); let mut env = HashMap::default(); @@ -2901,11 +2869,11 @@ def @main(%data: Tensor[(1, 2, 2, 16), float32], %bn_gamma: Tensor[(16), float32 env.insert(k.clone(), v.clone()); value_env.insert( k.clone(), - ndarray::ArrayD::::random_using( + ndarray::ArrayD::from_shape_vec( v.clone(), - Uniform::new(1f32, 2f32), // prevent NaN results - &mut tensor_rng, - ), + (0..v.iter().product::()).map(|x| x as f32).collect(), + ) + .unwrap(), ); } @@ -2915,830 +2883,111 @@ def @main(%data: Tensor[(1, 2, 2, 16), float32], %bn_gamma: Tensor[(16), float32 let id = egraph.add_expr(&expr); - let result = run_relay(&value_env, &shapes_vec, relay_to_run); + let result_output = run_relay(&value_env, &shapes_vec, relay); let code = codegen( &egraph, id, &HashMap::default(), - "relay_batchnorm", + "compiled_function", "", - &vec!["data", "bn_gamma", "bn_beta", "bn_mean", "bn_var"], + &vec!["data"], &generate_worklist_for_codegen(&egraph, id), true, ); - let main_code = format!( - " -#include -#include -#include \"{}\" - -{} -{} -{} -{} -{} -{} -{} -{} - -int main() {{ - relay_batchnorm(out, data, bn_gamma, bn_beta, bn_mean, bn_var); - - for (int i = 0; i < {}; i++) {{ - assert(fabs(((float*)result)[i] - ((float*)out)[i]) < 0.00001); - }} -}} -", - PathBuf::from_str( - format!( - "{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "c-files", - "relay-op-implementations.c" - ) - .as_str() - ) - .unwrap() - .to_string_lossy(), - c_assignment_string( - "", - "data", - DType::Fp32, - &value_env.get("data").unwrap().view() - ), - c_assignment_string( - "", - "bn_gamma", - DType::Fp32, - &value_env.get("bn_gamma").unwrap().view() - ), - c_assignment_string( - "", - "bn_beta", - DType::Fp32, - &value_env.get("bn_beta").unwrap().view() - ), - c_assignment_string( - "", - "bn_mean", - DType::Fp32, - &value_env.get("bn_mean").unwrap().view() - ), - c_assignment_string( - "", - "bn_var", - DType::Fp32, - &value_env.get("bn_var").unwrap().view() - ), - c_assignment_string("", "result", DType::Fp32, &result.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(result.shape()).view() - ), - code, - result.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-batchnorm-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-batchnorm-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .arg("-lm") - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - } - - #[test] - fn relay_op_softmax() { - let relay = r#" -#[version = "0.0.5"] -def @main(%data: Tensor[(1,100), float32]) -> Tensor[(1,100), float32] { - nn.softmax(%data) /* ty=Tensor[(1,10), float32] */ -} -"#; - - let module = tvm::ir::module::IRModule::parse("", relay).unwrap(); - - let (expr, shapes_vec) = crate::language::from_relay::from_relay( - &module, - true, - &vec![crate::language::RelayOperator::RelaySoftmax], - ); - - let mut env = HashMap::default(); - let mut value_env = HashMap::default(); - for (k, v) in &shapes_vec { - env.insert(k.clone(), v.clone()); - value_env.insert( - k.clone(), - ndarray::ArrayD::from_shape_vec( - v.clone(), - (0..v.iter().product::()).map(|x| x as f32).collect(), - ) - .unwrap(), - ); - } - - let mut egraph = EGraph::new(MyAnalysis { - name_to_shape: env.clone(), - }); - - let id = egraph.add_expr(&expr); - - let result_output = run_relay(&value_env, &shapes_vec, relay); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "relay_softmax", - "", - &vec!["data"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - let main_code = format!( - " -#include -#include -#include \"{}\" - -{} -{} -{} -{} - -int main() {{ - relay_softmax(out, data); - - for (int i = 0; i < {}; i++) {{ - assert(fabs(((float*)result)[i] - ((float*)out)[i]) < 0.00001); - }} -}} -", - PathBuf::from_str( - format!( - "{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "c-files", - "relay-op-implementations.c" - ) - .as_str() - ) - .unwrap() - .to_string_lossy(), - c_assignment_string( - "", - "data", - DType::Fp32, - &value_env.get("data").unwrap().view() - ), - c_assignment_string("", "result", DType::Fp32, &result_output.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(result_output.shape()).view() - ), - code, - result_output.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-softmax-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-softmax-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .arg("-lm") - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - } - - #[test] - fn relay_op_relu() { - let relay = r#" -#[version = "0.0.5"] -def @main(%x: Tensor[(1, 3, 3, 4), float32]) { - nn.relu(%x) -} -"#; - - // Random number generator for generating random tensors. - const SEED: u64 = 23; - let mut tensor_rng = SmallRng::seed_from_u64(SEED); - - let module = tvm::ir::module::IRModule::parse("", relay).unwrap(); - - let (expr, shapes_vec) = crate::language::from_relay::from_relay( - &module, - true, - &vec![crate::language::RelayOperator::RelayReLU], - ); - - let mut env = HashMap::default(); - let mut value_env = HashMap::default(); - for (k, v) in &shapes_vec { - env.insert(k.clone(), v.clone()); - value_env.insert( - k.clone(), - ndarray::ArrayD::::random_using( - v.clone(), - Uniform::new(-1f32, 1f32), - &mut tensor_rng, - ), - ); - } - - let mut egraph = EGraph::new(MyAnalysis { - name_to_shape: env.clone(), - }); - - let id = egraph.add_expr(&expr); - - let result = run_relay(&value_env, &shapes_vec, relay); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "relay_relu", - "", - &vec!["x"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include -#include \"{}\" - -{} -{} -{} -{} - -int main() {{ - relay_relu(out, x); - - for (int i = 0; i < {}; i++) {{ - assert(fabs(((float*)result)[i] - ((float*)out)[i]) < 0.00001); - }} -}} -", - PathBuf::from_str( - format!( - "{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "c-files", - "relay-op-implementations.c" - ) - .as_str() - ) - .unwrap() - .to_string_lossy(), - c_assignment_string("", "x", DType::Fp32, &value_env.get("x").unwrap().view()), - c_assignment_string("", "result", DType::Fp32, &result.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(result.shape()).view() - ), - code, - result.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-relu-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-relu-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .arg("-lm") - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - } - - #[test] - fn relay_op_maxpool2d_resnet_3x3() { - let relay = r#" -#[version = "0.0.5"] -def @main(%x: Tensor[(1, 112, 112, 64), float32]) -> Tensor[(1, 56, 56, 64), float32] { - nn.max_pool2d(%x, pool_size=[3, 3], strides=[2, 2], padding=[1, 1, 1, 1], layout="NHWC") /* ty=Tensor[(1, 56, 56, 64), float32] */ -} -"#; - const SEED: u64 = 23; - let mut tensor_rng = SmallRng::seed_from_u64(SEED); - - let module = tvm::ir::module::IRModule::parse("", relay.clone()).unwrap(); - - let (expr, shapes_vec) = crate::language::from_relay::from_relay( - &module, - true, - &vec![crate::language::RelayOperator::RelayMaxPool2D], - ); - - let mut env = HashMap::default(); - let mut value_env = HashMap::default(); - for (k, v) in &shapes_vec { - env.insert(k.clone(), v.clone()); - value_env.insert( - k.clone(), - ndarray::ArrayD::::random_using( - v.clone(), - Uniform::new(-1f32, 1f32), - &mut tensor_rng, - ), - ); - } - - let mut egraph = EGraph::new(MyAnalysis { - name_to_shape: env.clone(), - }); - - let id = egraph.add_expr(&expr); - - let result = run_relay(&value_env, &shapes_vec, relay); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "relay_maxpool", - "", - &vec!["x"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include -#include -#include \"{}\" - -{} -{} -{} -{} - -int main() {{ - relay_maxpool(out, x); - - for (int i = 0; i < {}; i++) {{ - assert(fabs(((float*)result)[i] - ((float*)out)[i]) < 0.00001); - }} -}} -", - PathBuf::from_str( - format!( - "{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "c-files", - "relay-op-implementations.c" - ) - .as_str() - ) - .unwrap() - .to_string_lossy(), - c_assignment_string("", "x", DType::Fp32, &value_env.get("x").unwrap().view()), - c_assignment_string("", "result", DType::Fp32, &result.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(result.shape()).view() - ), - code, - result.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-maxpool-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-maxpool-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .arg("-lm") - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - } - - #[test] - fn relay_op_batchflatten() { - let relay = r#" -#[version = "0.0.5"] -def @main(%x: Tensor[(1, 512, 1, 1), float32]) { - nn.batch_flatten(%x) -} -"#; - const SEED: u64 = 23; - let mut tensor_rng = SmallRng::seed_from_u64(SEED); - - let module = tvm::ir::module::IRModule::parse("", relay).unwrap(); - - let (expr, shapes_vec) = crate::language::from_relay::from_relay( - &module, - true, - &vec![crate::language::RelayOperator::RelayBatchFlatten], - ); - - let mut env = HashMap::default(); - let mut value_env = HashMap::default(); - for (k, v) in &shapes_vec { - env.insert(k.clone(), v.clone()); - value_env.insert( - k.clone(), - ndarray::ArrayD::::random_using( - v.clone(), - Uniform::new(-2f32, 2f32), - &mut tensor_rng, - ), - ); - } - - let mut egraph = EGraph::new(MyAnalysis { - name_to_shape: env.clone(), - }); - - let id = egraph.add_expr(&expr); - - let result = run_relay(&value_env, &shapes_vec, relay); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "relay_maxpool", - "", - &vec!["x"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include -#include -#include \"{}\" - -{} -{} -{} -{} - -int main() {{ - relay_maxpool(out, x); - - for (int i = 0; i < {}; i++) {{ - assert(fabs(((float*)result)[i] - ((float*)out)[i]) < 0.00001); - }} -}} -", - PathBuf::from_str( - format!( - "{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "c-files", - "relay-op-implementations.c" - ) - .as_str() - ) - .unwrap() - .to_string_lossy(), - c_assignment_string("", "x", DType::Fp32, &value_env.get("x").unwrap().view()), - c_assignment_string("", "result", DType::Fp32, &result.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(result.shape()).view() - ), - code, - result.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-maxpool-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-maxpool-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .arg("-lm") - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - } + let var_assignment_str = c_assignment_string( + "", + "data", + DType::Fp32, + &value_env.get("data").unwrap().view(), + ); + let var = ", data"; + relay_op_test!( + "relay_op_softmax", + var_assignment_str, + var, + "fabs(((float*)result)[i] - ((float*)out)[i]) < 0.00001", + code, + result_output + ); + } + #[test] + fn relay_op_relu() { + let relay = r#" +#[version = "0.0.5"] +def @main(%x: Tensor[(1, 3, 3, 4), float32]) { + nn.relu(%x) +} +"#; + relay_op_test!( + "relay_op_relu", + relay, + relay, + crate::language::RelayOperator::RelayReLU, + 23, + -1f32, + 1f32, + vec!["x"] + ); + } + #[test] + fn relay_op_maxpool2d_resnet_3x3() { + let relay = r#" +#[version = "0.0.5"] +def @main(%x: Tensor[(1, 112, 112, 64), float32]) -> Tensor[(1, 56, 56, 64), float32] { + nn.max_pool2d(%x, pool_size=[3, 3], strides=[2, 2], padding=[1, 1, 1, 1], layout="NHWC") /* ty=Tensor[(1, 56, 56, 64), float32] */ +} +"#; + relay_op_test!( + "relay_op_maxpool2d_resnet_3x3", + relay, + relay, + crate::language::RelayOperator::RelayMaxPool2D, + 23, + -1f32, + 1f32, + vec!["x"] + ); + } + #[test] + fn relay_op_batchflatten() { + let relay = r#" + #[version = "0.0.5"] + def @main(%x: Tensor[(1, 512, 1, 1), float32]) { + nn.batch_flatten(%x) + } + "#; + relay_op_test!( + "relay_op_batchflatten", + relay, + relay, + crate::language::RelayOperator::RelayBatchFlatten, + 23, + -2f32, + 2f32, + vec!["x"] + ); + } #[test] fn relay_op_globalavgpool2d() { let relay = r#" -#[version = "0.0.5"] -def @main(%x: Tensor[(1, 7, 7, 512), float32]) -> Tensor[(1, 1, 1, 512), float32] { - nn.global_avg_pool2d(%x, layout="NHWC") /* ty=Tensor[(1, 1, 1, 512), float32] */ -} -"#; - const SEED: u64 = 23; - let mut tensor_rng = SmallRng::seed_from_u64(SEED); - - let module = tvm::ir::module::IRModule::parse("", relay.clone()).unwrap(); - - let (expr, shapes_vec) = crate::language::from_relay::from_relay( - &module, - true, - &vec![crate::language::RelayOperator::RelayGlobalAvgPool2D], - ); - - let mut env = HashMap::default(); - let mut value_env = HashMap::default(); - for (k, v) in &shapes_vec { - env.insert(k.clone(), v.clone()); - value_env.insert( - k.clone(), - ndarray::ArrayD::::random_using( - v.clone(), - Uniform::new(-2f32, 2f32), - &mut tensor_rng, - ), - ); + #[version = "0.0.5"] + def @main(%x: Tensor[(1, 7, 7, 512), float32]) -> Tensor[(1, 1, 1, 512), float32] { + nn.global_avg_pool2d(%x, layout="NHWC") /* ty=Tensor[(1, 1, 1, 512), float32] */ } - - let mut egraph = EGraph::new(MyAnalysis { - name_to_shape: env.clone(), - }); - - let id = egraph.add_expr(&expr); - - let result = run_relay(&value_env, &shapes_vec, relay); - - let code = codegen( - &egraph, - id, - &HashMap::default(), - "relay_globalavgpool2d", - "", - &vec!["x"], - &generate_worklist_for_codegen(&egraph, id), - true, - ); - - let main_code = format!( - " -#include -#include -#include \"{}\" - -{} -{} -{} -{} - -int main() {{ - relay_globalavgpool2d(out, x); - - for (int i = 0; i < {}; i++) {{ - assert(fabs(((float*)result)[i] - ((float*)out)[i]) < 0.00001); - }} -}} -", - PathBuf::from_str( - format!( - "{}/{}/{}", - env!("CARGO_MANIFEST_DIR"), - "c-files", - "relay-op-implementations.c" - ) - .as_str() - ) - .unwrap() - .to_string_lossy(), - c_assignment_string("", "x", DType::Fp32, &value_env.get("x").unwrap().view()), - c_assignment_string("", "result", DType::Fp32, &result.view()), - c_assignment_string( - "", - "out", - DType::Fp32, - &ndarray::ArrayD::::zeros(result.shape()).view() - ), - code, - result.shape().iter().product::() - ); - - let main_c_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-globalavgpool2d-test-{}.c", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", main_c_filepath.to_string_lossy()); - - let binary_filepath = std::env::temp_dir().with_file_name(format!( - "relay-op-globalavgpool2d-test-{}", - std::time::SystemTime::now().elapsed().unwrap().as_nanos() - )); - println!("{}", binary_filepath.to_string_lossy()); - - File::create(&main_c_filepath) - .unwrap() - .write_all(main_code.as_bytes()) - .unwrap(); - - let result = Command::new("gcc") - .arg("-Werror") - .arg("-g") - .arg("-o") - .arg(&binary_filepath) - .arg(&main_c_filepath) - .arg("-lm") - .output() - .unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") - ); - - let result = Command::new(&binary_filepath).output().unwrap(); - - assert!( - result.status.success(), - "{}", - std::str::from_utf8(result.stderr.as_slice()) - .expect("Could not convert stderr to UTF8") + "#; + relay_op_test!( + "relay_op_globalavgpool2d", + relay, + relay, + crate::language::RelayOperator::RelayGlobalAvgPool2D, + 23, + -2f32, + 2f32, + vec!["x"] ); } - #[test] #[ignore = "unfinished test"] fn relay_op_leakyrelu() {