Merge pull request #188 from luminal-ai/visualization

Modules for visualization, serialized_egraph, egglog_utils, example for visualization

Author: jafioti

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "luminal"
 version = "0.2.0"
-edition = "2024"
+edition.workspace = true
 rust-version = "1.85"
 description = "Deep learning at the speed of light."
 license = "MIT OR Apache-2.0"
@@ -22,18 +22,23 @@ regex = "1.9.5"
 rustc-hash = "2.1.1"
 uuid = { version = "1.7.0", features = ["v4"] }
 as-any = "0.3.1"
-egg = "0.9.5"
 symbolic_expressions = "5.0.3"
 serde = { version = "1.0.202", features = ["derive"] }
 thread_local = "1.1.8"
 generational-box = "0.5.6"
 serde_json = "1.0.140"
+egg = "0.9.5"
 egglog = "1.0.0"
 egglog-ast = "1.0.0"
-egraph-serialize = "0.3.0"
+egraph-serialize = { version = "0.3.0", default-features = false, features = ["graphviz", "serde"]}
 tracing = "0.1.43"
 paste = "1.0.15"
 pretty-duration = "0.1.1"
+anyhow = "1.0"
+graphviz-rust = { version = "0.9", default-features = false}
+
+[workspace.package]
+edition = "2024"
 
 [dev-dependencies]
 candle-core = "0.9.1"
@@ -43,13 +48,13 @@ ordered-float = "5.1.0"
 [workspace]
 members = [
     "examples/llama",
-    #"examples/*",
+    "examples/visualization",
     "crates/luminal_nn",
     "crates/luminal_cuda",
     "crates/luminal_training",
     "docs/company",
 ]
 exclude = [
-    "examples/yolo_v8",
+
     "crates/luminal_cuda",
 ]

crates/luminal_cuda/src/block/ops.rs

@@ -4,8 +4,9 @@ use super::CustomState;
 use cudarc::driver::{CudaStream, DevicePtr};
 use itertools::Itertools;
 use luminal::{
-    graph::{extract_expr, extract_expr_list, SerializedEGraph},
+    graph::{extract_expr, extract_expr_list},
     prelude::ENodeId,
+    serialized_egraph::SerializedEGraph,
     shape::Expression,
     utils::{
         flatten_mul_strides, CStructBuilder, EgglogOp, LLIROp,

crates/luminal_cuda/src/kernel/ops.rs

@@ -6,9 +6,10 @@ use cudarc::{
 };
 use itertools::Itertools;
 use luminal::{
-    graph::{extract_dtype, extract_expr, extract_expr_list, SerializedEGraph},
+    graph::{extract_dtype, extract_expr, extract_expr_list},
     op::DType,
     prelude::ENodeId,
+    serialized_egraph::SerializedEGraph,
     shape::Expression,
     utils::{
         flatten_mul_strides, EgglogOp, LLIROp,

crates/luminal_nn/src/lib.rs

@@ -1,3 +1,5 @@
+#![allow(unused_imports)]
+
 mod activation;
 pub use activation::*;
 mod convolution;

examples/llama/Cargo.toml

@@ -12,7 +12,6 @@ luminal_cuda = { path = "../../crates/luminal_cuda" }
 itertools = "0.12.1"
 tokenizers = "0.15.2"
 tracing = "0.1.43"
-rustc-hash = "2.1.1"
 tracing-subscriber = {version="0.3", features=["env-filter"]}
 tracing-perfetto-sdk-layer = "0.13.0"
 tracing-perfetto-sdk-schema = "0.13.0"

examples/llama/src/main.rs

@@ -4,13 +4,13 @@ use itertools::Itertools;
 use luminal::{
     graph::{Graph, Runtime},
     op::DType,
+    prelude::FxHashMap,
 };
 use luminal_cuda::{
     block::IntoBlockOp,
     runtime::{record_exec_timings_to_file, CudaRuntime, CustomState},
 };
 use model::*;
-use rustc_hash::*;
 use std::{fs::File, io::Write, time::Duration};
 use tokenizers::Tokenizer;
 use tracing::{span, Level};

examples/visualization/.gitignore

@@ -0,0 +1,3 @@
+*.dot
+*.html
+*.svg

examples/visualization/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "visualization"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+
+[dependencies]
+anyhow = "1.0"
+egglog = "1.0"
+egglog-ast = "1.0.0"
+egraph-serialize = { version = "0.3.0", default-features = false, features = ["graphviz", "serde"]}
+itertools = "0.12.1"
+luminal = { path = "../.." }
+luminal_cuda = { path = "../../crates/luminal_cuda" }
+luminal_nn = { path = "../../crates/luminal_nn" }
+rustc-hash = "2.1"
+tokenizers = "0.15.2"
+tracing = "0.1.43"
+tracing-appender = "0.2.4"
+tracing-perfetto-sdk-layer = "0.13.0"
+tracing-perfetto-sdk-schema = "0.13.0"
+tracing-perfetto-sdk-sys = "0.13.0"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }

examples/visualization/README.md

@@ -0,0 +1,23 @@
+# Visualization in Luminal
+
+## Design Choices
+Luminal produces intermediate files rather than complete visualizations
+
+The two primary file types are: 
+- `.html` files
+- `.dot` files
+
+These files enable interactive viewing which is often necessary for making visualizations interpretable. 
+
+## VSCode Extensions
+We recommend the following extensions for VSCode users. 
+The integrated nature of these extensions makes viewing these files easy even on remote machines via ssh. 
+
+- `Live Preview` by microsoft. 
+- `Graphviz Interactive Preview` by tintinweb 
+
+## Example Provided 
+In the example program, as simple program is defined. 
+From this a HLIR graph is created and visualized. 
+A saturated EGraph is created and visualized. 
+Finally an LLIR graph is extracted and visualized. 

examples/visualization/src/main.rs

@@ -0,0 +1,86 @@
+use std::fs;
+
+use luminal::{
+    self,
+    graph::{hlir_to_egglog, Graph, Runtime},
+    prelude::*,
+    serialized_egraph::SerializedEGraph,
+    visualization::{ToDot, ToHtml},
+};
+use luminal_cuda::runtime::{CudaRuntime, CustomState};
+
+use egglog::{prelude::RustSpan, var, EGraph};
+use egglog_ast::span::Span;
+use rustc_hash::FxHashMap;
+
+fn main() {
+    // Create a new graph
+    let mut cx = Graph::new();
+
+    // Create input tensor using constant values
+
+    let (m, n, k) = (4096, 14336, 9192);
+
+    let a = cx.tensor((m, k));
+    let b = cx.tensor((k, n));
+
+    let _c = a.matmul(b);
+
+    let ctx = luminal_cuda::cudarc::driver::CudaContext::new(0).unwrap();
+    ctx.bind_to_thread().unwrap();
+    let _stream = ctx.default_stream();
+    let _custom_state: FxHashMap<String, CustomState> = FxHashMap::default();
+
+    println!("Visualizing HLIR");
+    fs::write("HLIR.dot", cx.graph.to_dot().unwrap()).unwrap();
+
+    println!("Building and Saturating EGraph");
+    cx.build_search_space::<CudaRuntime>();
+
+    let (program, root) = hlir_to_egglog(&cx);
+
+    let mut ops = <CudaRuntime as Runtime>::Ops::into_vec();
+    ops.extend(<luminal::op::Ops as IntoEgglogOp>::into_vec());
+
+    let mut egglog_obj: EGraph = egglog::EGraph::default();
+
+    // setup the rules and datatypes
+    egglog_obj
+        .parse_and_run_program(None, luminal::egglog_utils::BASE)
+        .unwrap();
+    egglog_obj
+        .parse_and_run_program(None, &luminal::egglog_utils::op_defs_string(&ops))
+        .unwrap();
+    egglog_obj
+        .parse_and_run_program(None, &luminal::egglog_utils::op_rewrites_string(&ops))
+        .unwrap();
+    egglog_obj
+        .parse_and_run_program(None, luminal::egglog_utils::BASE_CLEANUP)
+        .unwrap();
+    egglog_obj
+        .parse_and_run_program(None, &luminal::egglog_utils::op_cleanups_string(&ops))
+        .unwrap();
+
+    // load the program
+    egglog_obj.parse_and_run_program(None, &program).unwrap();
+
+    // run the graph
+    egglog_obj
+        .parse_and_run_program(None, luminal::egglog_utils::RUN_SCHEDULE)
+        .unwrap();
+
+    // EGraph Optimization Complete
+    println!("Visualizing EGraph");
+    // save the egraph visualizations
+    fs::write("egraph.html", egglog_obj.to_html().unwrap()).unwrap();
+    fs::write("egraph.dot", egglog_obj.to_dot().unwrap()).unwrap();
+
+    let (sort, value) = egglog_obj.eval_expr(&var!(root)).unwrap();
+    let s_egraph = SerializedEGraph::new(&egglog_obj, vec![(sort, value)]);
+    let llir_graphs = egglog_to_llir(&s_egraph, &ops, 100);
+
+    let example_llir_graph = llir_graphs.last().unwrap();
+
+    println!("Visualizing LLIR Graph");
+    fs::write("LLIR.dot", example_llir_graph.to_dot().unwrap()).unwrap();
+}