From 19fda7ddfb395487662813f37a45165e6d334625 Mon Sep 17 00:00:00 2001
From: Martijn Courteaux <courteauxmartijn@gmail.com>
Date: Tue, 22 Aug 2023 00:33:10 +0200
Subject: [PATCH] Add jump-buttons to get fro Stmt directly to Assembly

---
 src/StmtToViz.cpp | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/StmtToViz.cpp b/src/StmtToViz.cpp
index f4f06fda650f..f41a74424306 100644
--- a/src/StmtToViz.cpp
+++ b/src/StmtToViz.cpp
@@ -631,7 +631,9 @@ class HTMLCodePrinter : public IRVisitor {
         cost_model = std::move(cm);
     }
 
-    void print(const Module &m) {
+    void print(const Module &m, AssemblyInfo asm_info) {
+        assembly_info = std::move(asm_info);
+
         // Generate a unique ID for this module
         int id = gen_unique_id();
 
@@ -708,6 +710,7 @@ class HTMLCodePrinter : public IRVisitor {
 
     // Holds cost information for visualized program
     IRCostModel cost_model;
+    AssemblyInfo assembly_info;
 
     /* Private print functions to handle various IR types */
     void print(const Buffer<> &buf) {
@@ -928,6 +931,16 @@ class HTMLCodePrinter : public IRVisitor {
                << "</button>";
     }
 
+    // Prints a button to sync text with visualization
+    void print_assembly_button(const void *op) {
+        int asm_lno = assembly_info.get_asm_lno((uint64_t)op);
+        if (asm_lno != -1) {
+            stream << "<button class='icon-btn text-info sync-btn' onclick='scrollToAsm(\"" << asm_lno << "\")'>"
+                   << "  <i class='bi bi-arrow-right-square' title='Jump to Assembly'></i>"
+                   << "</button>";
+        }
+    }
+
     // CUDA kernels are embedded into modules as PTX assembly. This
     // routine pretty - prints that assembly format.
     void print_cuda_gpu_source_kernels(const std::string &str) {
@@ -1521,6 +1534,7 @@ class HTMLCodePrinter : public IRVisitor {
 
         // Add a button to jump to this producer/consumer in the viz
         print_visualization_button("prodcons-viz-" + std::to_string(id));
+        print_assembly_button(op);
 
         // Open code block to hold function body
         print_html_element("span", "matched", "{");
@@ -1583,6 +1597,7 @@ class HTMLCodePrinter : public IRVisitor {
 
         // Add a button to jump to this loop in the viz
         print_visualization_button("loop-viz-" + std::to_string(id));
+        print_assembly_button(op);
 
         // Open code block to hold function body
         print_html_element("span", "matched", "{");
@@ -2885,7 +2900,7 @@ class IRVisualizer {
     // Generate tab 1/3: Lowered IR code with syntax highlighting in HTML
     void generate_ir_tab(const Module &m) {
         stream << "<div id='ir-code-tab'>\n";
-        html_code_printer.print(m);
+        html_code_printer.print(m, asm_info);
         stream << "</div>\n";
     }