halide · maaz139 · Apr 12, 2023 · Sep 28, 2022 · Sep 28, 2022 · Oct 28, 2022
diff --git a/Makefile b/Makefile
@@ -561,6 +561,7 @@ SOURCE_FILES = \
   SplitTuples.cpp \
   StageStridedLoads.cpp \
   StmtToHtml.cpp \
+  StmtToViz.cpp \
   StorageFlattening.cpp \
   StorageFolding.cpp \
   StrictifyFloat.cpp \
@@ -585,6 +586,11 @@ SOURCE_FILES = \
    CodeGen_C_prologue \
    CodeGen_C_vectors
 
+HTML_TEMPLATE_FILES = \
+   StmtToViz_dependencies \
+   StmtToViz_javascript \
+   StmtToViz_stylesheet
+
 # The externally-visible header files that go into making Halide.h.
 # Don't include anything here that includes llvm headers.
 # Also *don't* include anything that's only used internally (eg SpirvIR.h).
@@ -734,6 +740,7 @@ HEADER_FILES = \
   SplitTuples.h \
   StageStridedLoads.h \
   StmtToHtml.h \
+  StmtToViz.h \
   StorageFlattening.h \
   StorageFolding.h \
   StrictifyFloat.h \
@@ -889,6 +896,7 @@ INITIAL_MODULES = $(RUNTIME_CPP_COMPONENTS:%=$(BUILD_DIR)/initmod.%_32.o) \
                   $(RUNTIME_CPP_COMPONENTS:%=$(BUILD_DIR)/initmod.%_64_debug.o) \
                   $(RUNTIME_EXPORTED_INCLUDES:$(INCLUDE_DIR)/%.h=$(BUILD_DIR)/initmod.%_h.o) \
                   $(C_TEMPLATE_FILES:%=$(BUILD_DIR)/c_template.%.o) \
+                  $(HTML_TEMPLATE_FILES:%=$(BUILD_DIR)/html_template.%.o) \
                   $(BUILD_DIR)/initmod.inlined_c.o \
                   $(RUNTIME_LL_COMPONENTS:%=$(BUILD_DIR)/initmod.%_ll.o) \
                   $(PTX_DEVICE_INITIAL_MODULES:libdevice.%.bc=$(BUILD_DIR)/initmod_ptx.%_ll.o)
@@ -1131,6 +1139,9 @@ $(BUILD_DIR)/initmod_ptx.%_ll.cpp: $(BIN_DIR)/binary2cpp $(SRC_DIR)/runtime/nvid
 $(BUILD_DIR)/c_template.%.cpp: $(BIN_DIR)/binary2cpp $(SRC_DIR)/%.template.cpp
 	./$(BIN_DIR)/binary2cpp halide_c_template_$* < $(SRC_DIR)/$*.template.cpp > $@
 
+$(BUILD_DIR)/html_template.%.cpp: $(BIN_DIR)/binary2cpp $(SRC_DIR)/irvisualizer/%.template.html
+	./$(BIN_DIR)/binary2cpp halide_html_template_$* < $(SRC_DIR)/irvisualizer/$*.template.html > $@
+
 $(BIN_DIR)/binary2cpp: $(ROOT_DIR)/tools/binary2cpp.cpp
 	@mkdir -p $(@D)
 	$(CXX) $< -o $@
@@ -1148,6 +1159,9 @@ $(BUILD_DIR)/initmod.%.o: $(BUILD_DIR)/initmod.%.cpp
 $(BUILD_DIR)/c_template.%.o: $(BUILD_DIR)/c_template.%.cpp
 	$(CXX) -c $< -o $@ -MMD -MP -MF $(BUILD_DIR)/$*.d -MT $(BUILD_DIR)/$*.o
 
+$(BUILD_DIR)/html_template.%.o: $(BUILD_DIR)/html_template.%.cpp
+	$(CXX) -c $< -o $@ -MMD -MP -MF $(BUILD_DIR)/$*.d -MT $(BUILD_DIR)/$*.o
+
 $(BUILD_DIR)/%.o: $(SRC_DIR)/%.cpp $(BUILD_DIR)/llvm_ok
 	@mkdir -p $(@D)
 	$(CXX) $(CXX_FLAGS) -c $< -o $@ -MMD -MP -MF $(BUILD_DIR)/$*.d -MT $(BUILD_DIR)/$*.o

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -149,6 +149,7 @@ set(HEADER_FILES
     SplitTuples.h
     StageStridedLoads.h
     StmtToHtml.h
+    StmtToViz.h
     StorageFlattening.h
     StorageFolding.h
     StrictifyFloat.h
@@ -329,6 +330,7 @@ set(SOURCE_FILES
     SplitTuples.cpp
     StageStridedLoads.cpp
     StmtToHtml.cpp
+    StmtToViz.cpp
     StorageFlattening.cpp
     StorageFolding.cpp
     StrictifyFloat.cpp
@@ -350,11 +352,17 @@ set(SOURCE_FILES
     WrapCalls.cpp
     )
 
-set(TEMPLATE_FILES
+set(C_TEMPLATE_FILES
     CodeGen_C_prologue
     CodeGen_C_vectors
     )
 
+set(HTML_TEMPLATE_FILES
+    StmtToViz_dependencies
+    StmtToViz_javascript
+    StmtToViz_stylesheet
+    )
+
 ##
 # Build and import the runtime.
 ##
@@ -366,7 +374,7 @@ add_subdirectory(runtime)
 ##
 
 add_library(Halide_c_templates OBJECT)
-foreach (f IN LISTS TEMPLATE_FILES)
+foreach (f IN LISTS C_TEMPLATE_FILES)
     set(SRC "$<SHELL_PATH:${CMAKE_CURRENT_SOURCE_DIR}/${f}.template.cpp>")
     set(DST "c_template.${f}.template.cpp")
 
@@ -377,6 +385,17 @@ foreach (f IN LISTS TEMPLATE_FILES)
     target_sources(Halide_c_templates PRIVATE ${DST})
 endforeach ()
 
+foreach (f IN LISTS HTML_TEMPLATE_FILES)
+    set(SRC "$<SHELL_PATH:${CMAKE_CURRENT_SOURCE_DIR}/irvisualizer/${f}.template.html>")
+    set(DST "html_template.${f}.template.cpp")
+
+    add_custom_command(OUTPUT "${DST}"
+                       COMMAND binary2cpp "halide_html_template_${f}" < "${SRC}" > "${DST}"
+                       DEPENDS "${SRC}" binary2cpp
+                       VERBATIM)
+    target_sources(Halide_c_templates PRIVATE ${DST})
+endforeach ()
+
 ##
 # Build the Halide mono-header.
 ##

diff --git a/src/CodeGen_LLVM.cpp b/src/CodeGen_LLVM.cpp
@@ -3660,11 +3660,13 @@ void CodeGen_LLVM::return_with_error_code(llvm::Value *error_code) {
 }
 
 void CodeGen_LLVM::visit(const ProducerConsumer *op) {
+    producer_consumer_id++;
+
     string name;
     if (op->is_producer) {
-        name = std::string("produce ") + op->name;
+        name = std::to_string(producer_consumer_id) + std::string("_produce_") + op->name;
     } else {
-        name = std::string("consume ") + op->name;
+        name = std::to_string(producer_consumer_id) + std::string("_consume_") + op->name;
     }
     BasicBlock *produce = BasicBlock::Create(*context, name, function);
     builder->CreateBr(produce);
@@ -3690,9 +3692,12 @@ void CodeGen_LLVM::visit(const For *op) {
         BasicBlock *preheader_bb = builder->GetInsertBlock();
 
         // Make a new basic block for the loop
-        BasicBlock *loop_bb = BasicBlock::Create(*context, std::string("for ") + op->name, function);
+        for_loop_id++;
+        BasicBlock *loop_bb = BasicBlock::Create(
+            *context, std::to_string(for_loop_id) + std::string("_for_") + op->name, function);
         // Create the block that comes after the loop
-        BasicBlock *after_bb = BasicBlock::Create(*context, std::string("end for ") + op->name, function);
+        BasicBlock *after_bb = BasicBlock::Create(
+            *context, std::to_string(for_loop_id) + std::string("_end_for_") + op->name, function);
 
         // If min < max, fall through to the loop bb
         Value *enter_condition = builder->CreateICmpSLT(min, max);

diff --git a/src/CodeGen_LLVM.h b/src/CodeGen_LLVM.h
@@ -665,6 +665,13 @@ class CodeGen_LLVM : public IRVisitor {
      */
     int effective_vscale;
 
+    /** Assign a unique ID to each producer-consumer and for-loop node. The IDs
+     * are printed as comments in assembly and used to link visualizations with
+     * the generated assembly code within `StmtToViz`
+     */
+    int producer_consumer_id = 0;
+    int for_loop_id = 0;
+
     /** Embed an instance of halide_filter_metadata_t in the code, using
      * the given name (by convention, this should be ${FUNCTIONNAME}_metadata)
      * as extern "C" linkage. Note that the return value is a function-returning-

diff --git a/src/Generator.cpp b/src/Generator.cpp
@@ -651,7 +651,7 @@ gengen
  -e  A comma separated list of files to emit. Accepted values are:
      [assembly, bitcode, c_header, c_source, cpp_stub, featurization,
       llvm_assembly, object, python_extension, pytorch_wrapper, registration,
-      schedule, static_library, stmt, stmt_html, compiler_log].
+      schedule, static_library, stmt, stmt_html, stmt_viz, compiler_log].
      If omitted, default value is [c_header, static_library, registration].
 
  -p  A comma-separated list of shared libraries that will be loaded before the
@@ -790,6 +790,13 @@ gengen
             output_types.insert(OutputFileType::registration);
             output_types.insert(OutputFileType::static_library);
         } else {
+            // if emit_flags contains "stmt_viz" but not "assembly", throw an error
+            bool has_stmt_viz = std::find(emit_flags.begin(), emit_flags.end(), "stmt_viz") != emit_flags.end();
+            bool has_assembly = std::find(emit_flags.begin(), emit_flags.end(), "assembly") != emit_flags.end();
+
+            user_assert(!has_stmt_viz || has_assembly)
+                << "Output flag `stmt_viz` requires the `assembly` flag to also be set.";
+
             // Build a reverse lookup table. Allow some legacy aliases on the command line,
             // to allow legacy build systems to work more easily.
             std::map<std::string, OutputFileType> output_name_to_enum = {

diff --git a/src/Module.cpp b/src/Module.cpp
@@ -19,6 +19,7 @@
 #include "Pipeline.h"
 #include "PythonExtensionGen.h"
 #include "StmtToHtml.h"
+#include "StmtToViz.h"
 
 namespace Halide {
 namespace Internal {
@@ -49,6 +50,7 @@ std::map<OutputFileType, const OutputInfo> get_output_info(const Target &target)
         {OutputFileType::static_library, {"static_library", is_windows_coff ? ".lib" : ".a", IsSingle}},
         {OutputFileType::stmt, {"stmt", ".stmt", IsMulti}},
         {OutputFileType::stmt_html, {"stmt_html", ".stmt.html", IsMulti}},
+        {OutputFileType::stmt_viz, {"stmt_viz", ".stmt.viz.html", IsMulti}},
     };
     return ext;
 }
@@ -551,6 +553,7 @@ void Module::compile(const std::map<OutputFileType, std::string> &output_files)
         std::map<OutputFileType, std::string> output_files_copy = output_files;
         output_files_copy.erase(OutputFileType::stmt);
         output_files_copy.erase(OutputFileType::stmt_html);
+        output_files_copy.erase(OutputFileType::stmt_viz);
         resolve_submodules().compile(output_files_copy);
         return;
     }
@@ -610,6 +613,11 @@ void Module::compile(const std::map<OutputFileType, std::string> &output_files)
             compile_llvm_module_to_llvm_assembly(*llvm_module, *out);
         }
     }
+
+    if (contains(output_files, OutputFileType::stmt_viz)) {
+        debug(1) << "Module.compile(): stmt_viz " << output_files.at(OutputFileType::stmt_viz) << "\n";
+        Internal::print_to_viz(output_files.at(OutputFileType::stmt_viz), *this);
+    }
     if (contains(output_files, OutputFileType::function_info_header)) {
         debug(1) << "Module.compile(): function_info_header " << output_files.at(OutputFileType::function_info_header) << "\n";
         std::ofstream file(output_files.at(OutputFileType::function_info_header));

diff --git a/src/Module.h b/src/Module.h
@@ -41,6 +41,7 @@ enum class OutputFileType {
     static_library,
     stmt,
     stmt_html,
+    stmt_viz,
 };
 
 /** Type of linkage a function in a lowered Halide module can have.