Tuples!

Cargo.toml

@@ -39,8 +39,8 @@ git = "https://github.com/gussmith23/rplex"
 # issue.
 # NOTE Keep glenside-evaluation in sync with this
 # If the versions get out of whack, we'll probably have some weird errors.
-rev = "ff5450f9cbe07d57e73594c5fbd68f3bd665c2dc"
-git = "https://github.com/gussmith23/tvm"
+rev = "4ed759f6b6cafbab707f95b57762913a1f57c021"
+git = "https://github.com/hypercubestart/incubator-tvm"
 optional = true
 
 [dependencies.egg]
@@ -53,4 +53,4 @@ features = ["approx"]
 
 [dependencies.serde]
 version = "1.0"
-features = ["derive"]
+features = ["derive"]

Dockerfile

@@ -27,10 +27,10 @@ ENV LLVM_CONFIG_PATH=/usr/lib/llvm-10/bin/llvm-config
 # Build TVM with Rust bindings
 # TODO(@gussmith23) Switch this to TVM mainline
 # once https://github.com/apache/incubator-tvm/pull/6563 is merged
-RUN cd /root && git clone https://github.com/gussmith23/tvm tvm --recursive
+RUN cd /root && git clone https://github.com/hypercubestart/incubator-tvm tvm --recursive
 WORKDIR /root/tvm
 RUN git fetch
-RUN git checkout 14a786d3885304e5964df397a50edd19a759f903
+RUN git checkout 4ed759f6b6cafbab707f95b57762913a1f57c021
 RUN git submodule sync && git submodule update
 RUN echo 'set(USE_LLVM $ENV{LLVM_CONFIG_PATH})' >> config.cmake
 RUN echo 'set(USE_RPC ON)' >> config.cmake

src/codegen.rs

@@ -306,7 +306,10 @@ pub fn find_vars(expr: &Expr, id: Id) -> Vec<String> {
                 find_vars_recursive_helper(set, expr, id);
             }
             // Box<[Id]>
-            Language::RelayOperatorCall(ids) | Language::List(ids) | Language::Shape(ids) => {
+            Language::RelayOperatorCall(ids)
+            | Language::List(ids)
+            | Language::Shape(ids)
+            | Language::ConstructTuple(ids) => {
                 for id in ids.iter() {
                     find_vars_recursive_helper(set, expr, *id);
                 }
@@ -324,7 +327,8 @@ pub fn find_vars(expr: &Expr, id: Id) -> Vec<String> {
             | &Language::ShapeInsertAxis(ids)
             | &Language::ShapeRemoveAxis(ids)
             | &Language::AccessShape(ids)
-            | &Language::AccessSqueeze(ids) => {
+            | &Language::AccessSqueeze(ids)
+            | &Language::TupleGetItem(ids) => {
                 for id in ids.iter() {
                     find_vars_recursive_helper(set, expr, *id);
                 }
@@ -411,7 +415,10 @@ pub fn generate_worklist_for_codegen(expr: &Expr, id: Id) -> Vec<Id> {
                 }
             }
             // Box<[Id]>
-            Language::RelayOperatorCall(ids) | Language::Shape(ids) | Language::List(ids) => {
+            Language::RelayOperatorCall(ids)
+            | Language::Shape(ids)
+            | Language::List(ids)
+            | Language::ConstructTuple(ids) => {
                 for id in ids.iter() {
                     helper(worklist, expr, *id);
                 }
@@ -423,7 +430,8 @@ pub fn generate_worklist_for_codegen(expr: &Expr, id: Id) -> Vec<Id> {
             | &Language::AccessReshape(ids)
             | &Language::ShapeInsertAxis(ids)
             | &Language::ShapeRemoveAxis(ids)
-            | &Language::AccessSqueeze(ids) => {
+            | &Language::AccessSqueeze(ids)
+            | &Language::TupleGetItem(ids) => {
                 for id in ids.iter() {
                     helper(worklist, expr, *id);
                 }
@@ -1043,6 +1051,12 @@ add_with_broadcasting((float*) {out}, (float*) {X}, (float*) {Y}, (int*)  {out_s
 
                     Some(add_out)
                 }
+                RelayOperator::RelayLeakyReLU => todo!(),
+                RelayOperator::RelaySigmoid => todo!(),
+                RelayOperator::RelayAvgPool2D => todo!(),
+                RelayOperator::RelayUpSampling => todo!(),
+                RelayOperator::RelayMaximum => todo!(),
+                RelayOperator::RelayMinimum => todo!(),
             }
         }
         &Language::AccessWindows([access_id, filters_shape_id, stride_shape_id]) => {
@@ -1730,6 +1744,8 @@ if (i{i} < {dim_len}) {{
         | Language::RelayOperator(_) => None,
 
         &Language::Literal(_)
+        | &Language::ConstructTuple(_)
+        | &Language::TupleGetItem(_)
         | &Language::SystolicArrayConv2dIm2colNchwOihwWithBlocking(_)
         | &Language::SystolicArrayConv2dIm2colNhwcHwioWithBlocking(_)
         | &Language::SystolicArrayConv2dNchwOihwWithBlocking(_)
@@ -1761,7 +1777,7 @@ mod tests {
     use ndarray::{SliceInfo, SliceOrIndex};
     use ndarray_npy::{read_npy, write_npy};
     use ndarray_rand::{rand_distr::Uniform, RandomExt};
-    use rand::{rngs::SmallRng, Rng, SeedableRng};
+    use rand::{rngs::SmallRng, SeedableRng};
     use std::fs::File;
     use std::io::Write;
     use std::iter::FromIterator;
@@ -1795,7 +1811,9 @@ mod tests {
 
         let mut cmd = Command::new("python3");
         cmd.arg(script_filepath);
+        cmd.arg("--npy_out_filepath");
         cmd.arg(&output_filepath);
+        cmd.arg("--npy_arg_filepath");
         cmd.stdin(std::process::Stdio::piped())
             .stdout(std::process::Stdio::piped())
             .stderr(std::process::Stdio::piped());
@@ -1836,6 +1854,90 @@ mod tests {
         relay_output
     }
 
+    fn run_relay_tuple_out(
+        env: &HashMap<String, ArrayD<f32>>,
+        shapes_vec: &Vec<(String, Vec<usize>)>,
+        relay_str: &str,
+        outputs: usize,
+    ) -> Vec<ArrayD<f32>> {
+        let script_filepath = format!(
+            "{}/src/language/from_relay/run_relay.py",
+            env!("CARGO_MANIFEST_DIR")
+        );
+
+        let mut cmd = Command::new("python3");
+        cmd.arg(script_filepath);
+        cmd.arg("--npy_out_filepath");
+        let mut output_paths = Vec::new();
+        for _ in 0..outputs {
+            // https://www.reddit.com/r/rust/comments/38jhva/piping_string_to_child_process_stdin/crvlqcd/?utm_source=reddit&utm_medium=web2x&context=3
+            // Output filename
+            // TODO(@gussmith23) Do we want this RNG to use SEED?
+            // I initially attempted to do this, but was running into issues
+            // (I think the same filename kept being generated b/c I wasn't
+            // using the RNG carefully...but maybe there's also something
+            // wrong w/ how I'm reading files!)
+            let output_filepath = std::env::temp_dir().with_file_name(format!(
+                "output-{}.npy",
+                rand::thread_rng()
+                    .sample_iter(&rand::distributions::Alphanumeric)
+                    .take(30)
+                    .collect::<String>()
+            ));
+            cmd.arg(&output_filepath);
+            output_paths.push(
+                output_filepath
+                    .clone()
+                    .into_os_string()
+                    .into_string()
+                    .unwrap(),
+            );
+        }
+        cmd.arg("--npy_arg_filepath");
+        cmd.stdin(std::process::Stdio::piped())
+            .stdout(std::process::Stdio::piped())
+            .stderr(std::process::Stdio::piped());
+        for (name, _) in shapes_vec.iter() {
+            let value = env.get(name).unwrap();
+            // TODO(@gussmith23) output type assumption
+            let filepath = std::env::temp_dir().with_file_name(format!(
+                "arg-{}.npy",
+                rand::thread_rng()
+                    .sample_iter(&rand::distributions::Alphanumeric)
+                    .take(30)
+                    .collect::<String>()
+            ));
+            write_npy(&filepath, value).unwrap();
+            cmd.arg(filepath);
+        }
+
+        let mut proc = cmd.spawn().ok().expect("Failed to spawn process");
+        proc.stdin
+            .as_mut()
+            .unwrap()
+            .write_all(relay_str.as_bytes())
+            .unwrap();
+        let output = proc.wait_with_output().unwrap();
+        // Check that it ran.
+        assert!(
+            output.status.success(),
+            "Running Relay code failed with code {:?}.\nstdout:\n{}\nstderr:\n{}",
+            output.status.code(),
+            std::str::from_utf8(output.stdout.as_slice())
+                .expect("Could not convert stderr to UTF8"),
+            std::str::from_utf8(output.stderr.as_slice())
+                .expect("Could not convert stderr to UTF8")
+        );
+
+        // TODO(@gussmith23) output type assumption
+        let mut relay_outputs = Vec::new();
+        for output_filepath in output_paths.iter() {
+            let relay_output: ndarray::ArrayD<f32> = read_npy(output_filepath).unwrap();
+            relay_outputs.push(relay_output);
+        }
+        relay_outputs
+    }
+
     #[test]
     fn transpose() {
         let shape = vec![1, 20, 300, 3];
@@ -4042,4 +4144,107 @@ int main() {{
                 .expect("Could not convert stderr to UTF8")
         );
     }
+
+    #[test]
+    #[should_panic(expected = "unfinished test")]
+    fn relay_model_yolov3() {
+        // Generate yolov3 with directions from:
+        // https://tvm.apache.org/docs/tutorials/frontend/from_darknet.html
+        let filename = PathBuf::from(format!(
+            "{}/models/yolov3.relay",
+            env!("CARGO_MANIFEST_DIR")
+        ));
+        let relay = std::fs::read_to_string(&filename).unwrap();
+        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::RelayBatchNormInference,
+                crate::language::RelayOperator::RelaySoftmax,
+                crate::language::RelayOperator::RelayLeakyReLU,
+                crate::language::RelayOperator::RelayBiasAdd,
+                crate::language::RelayOperator::RelayAdd,
+                crate::language::RelayOperator::RelaySigmoid,
+                crate::language::RelayOperator::RelayUpSampling,
+            ],
+        );
+        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::<f32>::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_tuple_out(&value_env, &shapes_vec, &relay, 12);
+        todo!("unfinished test")
+    }
+
+    #[test]
+    #[should_panic(expected = "unfinished test")]
+    fn relay_model_efficientnet_lite4_11() {
+        // efficientnet onnx model source: https://github.com/onnx/models/blob/master/vision/classification/efficientnet-lite4/model/efficientnet-lite4-11.onnx
+        // imported into relay
+        let filename = PathBuf::from(format!(
+            "{}/models/efficientnet-lite4-11.relay",
+            env!("CARGO_MANIFEST_DIR")
+        ));
+        let relay = std::fs::read_to_string(&filename).unwrap();
+        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::RelayBatchNormInference,
+                crate::language::RelayOperator::RelaySoftmax,
+                crate::language::RelayOperator::RelayLeakyReLU,
+                crate::language::RelayOperator::RelayBiasAdd,
+                crate::language::RelayOperator::RelayAdd,
+                crate::language::RelayOperator::RelayMaximum,
+                crate::language::RelayOperator::RelayMinimum,
+                crate::language::RelayOperator::RelayAvgPool2D,
+            ],
+        );
+        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::<f32>::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);
+        todo!("unfinished test")
+    }
 }

src/extraction/mod.rs

@@ -69,6 +69,8 @@ impl egg::CostFunction<Language> for MonolithicCostFunction<'_> {
             | Language::SystolicArray(_)
             | Language::SystolicArrayWithBlocking(_)
             | Language::Usize(_)
+            | Language::ConstructTuple(_)
+            | Language::TupleGetItem(_)
             | Language::AccessSlice(_)
             | Language::AccessConcatenate(_)
             | Language::AccessPad(_)
@@ -173,6 +175,8 @@ impl CostFunction<Language> for SimpleCostFunction {
             Language::SystolicArrayConv2dNhwcHwioWithBlocking(_) => todo!(),
             Language::SystolicArrayConv2dIm2colNchwOihwWithBlocking(_) => todo!(),
             Language::SystolicArrayConv2dIm2colNhwcHwioWithBlocking(_) => todo!(),
+            Language::ConstructTuple(_) => todo!(),
+            Language::TupleGetItem(_) => todo!(),
 
             // Cannot extract compute: compute must be lowered to an atom.
             Compute(_) => std::usize::MAX,