Add codegen for multi-loops.

TiledTensor · May 30, 2024 · 3a88bc4 · 3a88bc4
1 parent 98bd470
commit 3a88bc4
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 18 deletions.
diff --git a/examples/loop/multiloops.rs b/examples/loop/multiloops.rs
@@ -2,7 +2,7 @@ use std::rc::Rc;
 
 use thriller_core::{
     initialize, AccessMap, AccessMatrix, AccessOffset, AttachedEdge, BlockType, Buffer,
-    IterationBound, IterationVar, MemoryLevel, ThrillerBlock, ThrillerGraph,
+    IterationBound, IterationVar, MemoryLevel, Task, ThrillerBlock, ThrillerGraph,
 };
 
 fn main() {
@@ -45,13 +45,14 @@ fn main() {
     let s_b = Rc::new(Buffer::new("sB"));
     let r_b = Rc::new(Buffer::new("rB"));
 
-    let in_edge0 = AttachedEdge::new(s_a, r_a.clone(), Some(access_map1.clone()));
-    let in_edge1 = AttachedEdge::new(s_b, r_b.clone(), Some(access_map2.clone()));
+    let in_edge0 = AttachedEdge::new(s_a.clone(), r_a.clone(), Some(access_map1.clone()));
+    let in_edge1 = AttachedEdge::new(s_b.clone(), r_b.clone(), Some(access_map2.clone()));
+    let in_edge2 = AttachedEdge::new(s_a.clone(), r_b.clone(), Some(access_map1.clone()));
 
     let subgraph = ThrillerGraph::new(MemoryLevel::Register);
 
     let mut block = ThrillerBlock::new(
-        vec![Rc::new(in_edge0), Rc::new(in_edge1)],
+        vec![Rc::new(in_edge0), Rc::new(in_edge1), Rc::new(in_edge2)],
         vec![],
         MemoryLevel::Register,
         Rc::new(subgraph),
@@ -60,7 +61,7 @@ fn main() {
 
     block.merge_loops();
 
-    // let code = block.emit().unwrap();
+    let code = block.emit().unwrap();
 
-    // println!("{}", code);
+    println!("{}", code);
 }
diff --git a/thriller-core/src/dataflow/block.rs b/thriller-core/src/dataflow/block.rs
@@ -84,24 +84,16 @@ impl ThrillerBlock {
         let mut code = String::new();
 
         for edge in self.inputs.iter() {
-            if let Some(access) = edge.get_access() {
+            if edge.get_access().is_some() {
                 // TODO: Add access pattern support for load operation.
-                // let load = |access_map: &AccessMap| -> ThrillerResult<String> {
-                //     self.gen_load(access_map, edge)
-                // };
-                // code += access.gen_loop_access(&[load])?.as_str();
-                code += self.gen_load(access, edge)?.as_str();
+                code += self.gen_load(edge)?.as_str();
             }
         }
         Ok(code)
     }
 
     /// Generate load code for the block inputs.
-    pub(crate) fn gen_load(
-        &self,
-        _access_map: &AccessMap,
-        edge: &Rc<AttachedEdge>,
-    ) -> ThrillerResult<String> {
+    pub(crate) fn gen_load(&self, edge: &Rc<AttachedEdge>) -> ThrillerResult<String> {
         // TODO: This is not a final version of the load code generation. It is just a pseudocode representation of the formalized data flow.
         let mut code = String::new();
         // Generate load inputs.
@@ -258,7 +250,23 @@ impl ThrillerBlock {
                 let code = self.subgraph.emit()?;
                 Ok(code)
             } else {
-                unimplemented!();
+                // unimplemented!();
+                let mut code = String::new();
+                for group in self.loop_groups.iter() {
+                    let edges = &group.edges;
+                    let mut inner_code = String::new();
+
+                    for edge in edges.iter() {
+                        inner_code += &self.gen_load(edge)?;
+                    }
+
+                    let access_map = group.edges[0].get_access().as_ref().unwrap();
+                    code += access_map.gen_loop_access(inner_code)?.as_str();
+                }
+
+                // TODO: Add codegen for subgraph and split subgraph into different loops.
+
+                Ok(code)
             }
         }
     }