feat(core): Add Memory Access codegen based on AccessMap. (#15)

* Remove gen_load method.

* Remove BlockType.

* fix AccessMap.

* Add AttachedEdge comments.

* Update docs.

* fix compilation errors.

* fix AccessMap field.

* Add emit_access for AccessMap.

* Add memory access codegen for AttachedEdge.

* Update docs.

* fix gemm codegen.

* fix cargo clippy

README.md

@@ -1,10 +1,18 @@
 # ThrillerFlow
 
 ## Introduction
-**ThrillerFlow** is a **Dataflow Analysis** and **Codegen** Framework written in Rust. It has the following features:
-- Dataflow Analysis based on **Memory Hierarchy Graph**.
-- Loop Analysis based on **Polyhedral Model**.
-- Codegen based on **[TiledCUDA](https://github.com/TiledTensor/TiledCUDA)(An Efficient Kernel Template Library written in CuTe)**.
+**ThrillerFlow** is a **Dataflow Analysis** and **Codegen** Framework written in Rust. 
+
+In **ThrillerFlow**, we introduce a nested multi-dimendional dataflow
+graph called **Extended task Dependence Graph(ETDG)**, a unified intermediate
+representation that preserves a holistic view of parallelism and
+dependency across different control and data nested levels on the code.
+To facilitate code analysis and the later low-level code generation,
+an ETDG concisely encodes complex control structures and precisely
+represents the iteration-level data dependencies with and acyclic graph.
+For a clear exposition, ETDG borrows the concepts from the reduced dependence
+graph used classical compilers with the **Static Control Program(SCoP)** modeling
+employed in polyhedral compilers.
 
 ## Quick Start
 

examples/graph/topo_sort.rs

@@ -3,7 +3,7 @@ use std::{cell::RefCell, rc::Rc};
 
 use thriller_core::{
     initialize, AccessMap, AccessMatrix, AccessOffset, Gemm, IterationBound, IterationVar,
-    MemoryLevel, ThrillerEdge, ThrillerGraph, ThrillerNode, ThrillerNodeInner,
+    ThrillerEdge, ThrillerGraph, ThrillerNode, ThrillerNodeInner,
 };
 
 use thriller_utils::BufBuilder;
@@ -31,7 +31,7 @@ fn main() {
 
     let access_map = Rc::new(access_map);
 
-    let mut subgraph = ThrillerGraph::new(MemoryLevel::Register);
+    let mut subgraph = ThrillerGraph::new();
 
     let r_a_node = Rc::new(RefCell::new(ThrillerNode::new(ThrillerNodeInner::Buffer(
         r_a.clone(),

thriller-bindings/examples/gemm/gemm_g2r.py

@@ -46,7 +46,7 @@
     print(gC)
 
     MemoryLevel = pythriller.PyMemoryLevel.Register
-    RegGraph = pythriller.PyGraph(MemoryLevel)
+    RegGraph = pythriller.PyGraph()
 
     NodeA = pythriller.PyNode(rA)
     NodeB = pythriller.PyNode(rB)
@@ -57,11 +57,9 @@
     LoopIter = pythriller.IterationVar('i', (0, 4))
 
     access_dims = [1]
-    access_map = [[1]]
-    access_offset = [0]
 
     AccessMap = pythriller.AccessMap(
-        access_dims, access_map, access_offset, [LoopIter])
+        access_dims, [[[1]], [[0]]], [[0], [10]], [LoopIter])
 
     EdgeA_Gemm = pythriller.PyEdge(NodeA, GemmNode)
     EdgeB_GEMM = pythriller.PyEdge(NodeB, GemmNode)
@@ -77,10 +75,8 @@
     StoreRegToGlobalEdgeC = pythriller.AttachedEdge(acc, gC, AccessMap)
     G2RBlockMemLevel = pythriller.PyMemoryLevel.Register
 
-    G2RBlockType = pythriller.BlockType.Reduce
-
     GlobalToRegBlock = pythriller.Block(
-        [LoadGlobalToRegEdgeA, LoadGlobalToRegEdgeB], [StoreRegToGlobalEdgeC], RegGraph, G2RBlockType, [LoopIter])
+        [LoadGlobalToRegEdgeA, LoadGlobalToRegEdgeB], [StoreRegToGlobalEdgeC], RegGraph, [LoopIter])
 
     code = GlobalToRegBlock.codegen()
 

thriller-bindings/pythriller/__init__.py

@@ -1 +1 @@
-from .context import initialize_thriller_flow, PyLayout, PyBufType, PyBuffer, PyGraph, PyNode, PyEdge, PyMemoryLevel, Gemm, AttachedEdge, Block, BlockType, IterationVar, AccessMap
+from .context import initialize_thriller_flow, PyLayout, PyBufType, PyBuffer, PyGraph, PyNode, PyEdge, PyMemoryLevel, Gemm, AttachedEdge, Block, IterationVar, AccessMap

thriller-bindings/src/access.rs

@@ -28,30 +28,43 @@ impl PyAccessMap {
             .map(|v| v.extract::<PyRef<PyIterationVar>>().unwrap().0.clone())
             .collect::<Vec<_>>();
 
-        let access = access
+        let accesses = access
             .into_iter()
             .map(|a| {
                 a.extract::<Bound<PyList>>()
                     .unwrap()
                     .into_iter()
-                    .map(|i| i.extract::<usize>().unwrap())
+                    .map(|i| {
+                        i.extract::<Bound<PyList>>()
+                            .unwrap()
+                            .into_iter()
+                            .map(|j| j.extract::<usize>().unwrap())
+                            .collect::<Vec<_>>()
+                    })
                     .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>();
 
-        let access = AccessMatrix(access);
+        let access_matrixs = accesses.into_iter().map(AccessMatrix).collect::<Vec<_>>();
 
-        let offset = offset
+        let offsetes = offset
             .into_iter()
-            .map(|o| o.extract::<usize>().unwrap())
+            .map(|o| {
+                o.extract::<Bound<PyList>>()
+                    .unwrap()
+                    .into_iter()
+                    .map(|i| i.extract::<usize>().unwrap())
+                    .collect::<Vec<_>>()
+            })
             .collect::<Vec<_>>();
-        let offset = AccessOffset(offset);
+
+        let access_offsets = offsetes.into_iter().map(AccessOffset).collect::<Vec<_>>();
 
         let mut map = AccessMap::new(dims.len(), dims);
 
         map.add_iter_vars(vars);
-        map.add_access_matrix(access);
-        map.add_access_offset(offset);
+        map.add_access_matrixs(access_matrixs);
+        map.add_access_offsets(access_offsets);
 
         PyAccessMap(Rc::new(map))
     }

thriller-bindings/src/block.rs

@@ -1,17 +1,11 @@
 use std::rc::Rc;
 
-use thriller_core::{AttachedEdge, BlockType, Task, ThrillerBlock};
+use thriller_core::{AttachedEdge, Task, ThrillerBlock};
 
 use pyo3::{prelude::*, types::PyList};
 
 use crate::{access::PyAccessMap, buffer::PyBuffer, graph::PyGraph, var::PyIterationVar};
 
-#[pyclass(module = "block", name = "BlockType")]
-pub enum PyBlockType {
-    Reduce,
-    Map,
-}
-
 #[pyclass(unsendable, module = "block", name = "Block")]
 pub struct PyBlock(pub ThrillerBlock);
 
@@ -25,14 +19,8 @@ impl PyBlock {
         inputs: &Bound<PyList>,
         outputs: &Bound<PyList>,
         subgraph: PyRef<PyGraph>,
-        block_type: PyRef<PyBlockType>,
         ivars: &Bound<PyList>,
     ) -> PyResult<Self> {
-        let block_type = match *block_type {
-            PyBlockType::Reduce => BlockType::Reduce,
-            PyBlockType::Map => BlockType::Map,
-        };
-
         let inputs = inputs
             .into_iter()
             .map(|edge| {
@@ -62,7 +50,7 @@ impl PyBlock {
 
         let subgraph = Rc::clone(&subgraph.0);
 
-        let block = ThrillerBlock::new(inputs, outputs, subgraph, block_type, ivars);
+        let block = ThrillerBlock::new(inputs, outputs, subgraph, ivars);
 
         Ok(PyBlock(block))
     }
@@ -81,6 +69,6 @@ impl PyAttachedEdge {
         let src = Rc::clone(&src.0);
         let dst = Rc::clone(&dst.0);
         let map = Rc::clone(&map.0);
-        PyAttachedEdge(Rc::new(AttachedEdge::new(src, dst, Some(map))))
+        PyAttachedEdge(Rc::new(AttachedEdge::new(src, dst, map)))
     }
 }

thriller-bindings/src/graph.rs

@@ -2,8 +2,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyList;
 
 use thriller_core::{
-    AccessMap, Gemm, MemoryLevel, Task, ThrillerEdge, ThrillerGraph, ThrillerNode,
-    ThrillerNodeInner,
+    AccessMap, Gemm, Task, ThrillerEdge, ThrillerGraph, ThrillerNode, ThrillerNodeInner,
 };
 
 use crate::buffer::PyBuffer;
@@ -23,14 +22,8 @@ pub struct PyGraph(pub Rc<RefCell<ThrillerGraph>>);
 #[pymethods]
 impl PyGraph {
     #[new]
-    fn new(mem_level: &PyMemoryLevel) -> PyGraph {
-        let mem_level = match mem_level {
-            PyMemoryLevel::Register => MemoryLevel::Register,
-            PyMemoryLevel::Shared => MemoryLevel::Shared,
-            PyMemoryLevel::Global => MemoryLevel::Global,
-        };
-
-        PyGraph(Rc::new(RefCell::new(ThrillerGraph::new(mem_level))))
+    fn empty() -> PyGraph {
+        PyGraph(Rc::new(RefCell::new(ThrillerGraph::new())))
     }
 
     fn add_nodes(&mut self, nodes: &Bound<'_, PyList>) -> PyResult<()> {

thriller-bindings/src/lib.rs

@@ -1,7 +1,7 @@
 use access::PyAccessMap;
 use pyo3::prelude::*;
 
-use block::{PyAttachedEdge, PyBlock, PyBlockType};
+use block::{PyAttachedEdge, PyBlock};
 use buffer::{PyBufType, PyBuffer, PyLayout};
 use graph::{PyEdge, PyGraph, PyMemoryLevel, PyNode};
 use op::PyGemm;
@@ -40,7 +40,6 @@ fn thriller_flow(m: &Bound<'_, PyModule>) -> PyResult<()> {
 
     m.add_class::<PyBlock>()?;
     m.add_class::<PyAttachedEdge>()?;
-    m.add_class::<PyBlockType>()?;
 
     m.add_class::<PyIterationVar>()?;