feat: kernels to file from einsumtree

Author: RivinHD

CMakeLists.txt

@@ -311,6 +311,20 @@ if(OpenMP_CXX_FOUND)
     target_link_libraries(tests PRIVATE OpenMP::OpenMP_CXX)
 endif()
 
+# tests sanitized
+add_executable(tests_sanitized "${SOURCE_FILEPATHS}" "${TEST_FILEPATHS}")
+if(SAVE_JITS_TO_FILE)
+    target_compile_definitions(tests_sanitized PUBLIC SAVE_JITS_TO_FILE)
+endif(SAVE_JITS_TO_FILE)
+
+target_compile_options(tests_sanitized PRIVATE -g -fsanitize=float-divide-by-zero -fsanitize=bounds -fsanitize=address -fsanitize=undefined -fno-omit-frame-pointer)
+target_link_options(tests_sanitized PRIVATE -g -fsanitize=address -fsanitize=undefined)
+target_link_libraries(tests_sanitized PRIVATE Catch2::Catch2WithMain)
+
+if(OpenMP_CXX_FOUND)
+    target_link_libraries(tests_sanitized PRIVATE OpenMP::OpenMP_CXX)
+endif()
+
 # benchmarks
 add_executable(benchmarks "${SOURCE_FILEPATHS}" "${BENCH_FILEPATHS}")
 

src/main/Brgemm.cpp

@@ -47,6 +47,11 @@ mini_jit::Brgemm::kernel_t mini_jit::Brgemm::get_kernel() const
   return kernel;
 }
 
+void mini_jit::Brgemm::write_kernel_to_file(const char *path) const
+{
+  native_kernel.write(path);
+}
+
 void mini_jit::Brgemm::fill_with_matmuls_no_batch_dim_column_major_fp32(uint32_t m, uint32_t n, uint32_t k)
 {
   // Always sort from the specific to the more general case

src/main/Brgemm.h

@@ -66,6 +66,13 @@ class mini_jit::Brgemm
    **/
   kernel_t get_kernel() const;
 
+  /**
+   * @brief Writes the current kernel into a file.
+   *
+   * @param path The file to write the kernel to.
+   */
+  void write_kernel_to_file(const char *path) const;
+
 private:
   kernel_t kernel = nullptr;
   mini_jit::Kernel native_kernel;

src/main/EinsumTree.cpp

@@ -159,10 +159,16 @@ mini_jit::EinsumTree::~EinsumTree()
 
 void mini_jit::EinsumTree::delete_tree(EinsumNode *node)
 {
-  if (!node)
+  if (node == nullptr)
+  {
     return;
+  }
+
   delete_tree(node->left);
   delete_tree(node->right);
+  node->left = nullptr;
+  node->right = nullptr;
+
   if (node->type != NodeType::Leaf && node->tensor != nullptr && node != get_root())
   {
     delete[] node->tensor;
@@ -224,6 +230,16 @@ std::string mini_jit::EinsumTree::EinsumNode::to_string() const
   return mini_jit::EinsumTree::EinsumNode::_to_string(0, "", "");
 }
 
+std::string mini_jit::EinsumTree::EinsumNode::name() const
+{
+  std::string output = std::format("{}", output_dim_ids[0]);
+  for (auto iDim = output_dim_ids.begin() + 1; iDim != output_dim_ids.end(); iDim++)
+  {
+    output += std::format("_{}", *iDim);
+  }
+  return output;
+}
+
 mini_jit::TensorConfig mini_jit::EinsumTree::lower_node(const EinsumNode *node)
 {
   // Node has two children -> contraction
@@ -477,6 +493,9 @@ mini_jit::EinsumTree::ErrorExecute mini_jit::EinsumTree::execute_node(const std:
       return error;
     }
 
+#ifdef SAVE_JITS_TO_FILE
+    tensor_op.write_kernel_to_file(node->name());
+#endif  // SAVE_JITS_TO_FILE
     tensor_op.execute(node->left->tensor, nullptr, node->tensor);
   }
   else if (node->type == NodeType::Contraction)
@@ -517,6 +536,9 @@ mini_jit::EinsumTree::ErrorExecute mini_jit::EinsumTree::execute_node(const std:
       return error;
     }
 
+#ifdef SAVE_JITS_TO_FILE
+    tensor_op.write_kernel_to_file(node->name());
+#endif  // SAVE_JITS_TO_FILE
     tensor_op.execute(node->left->tensor, node->right->tensor, node->tensor);
   }
   else

src/main/EinsumTree.h

@@ -76,6 +76,11 @@ namespace mini_jit
        */
       std::string to_string() const;
 
+      /**
+       * @brief Gets the string representation of the dim ids of the node.
+       */
+      std::string name() const;
+
       /**
        * Get the size of the tensor represented by this node.
        *

src/main/TensorOperation.cpp

@@ -1,6 +1,7 @@
 #include "TensorOperation.h"
 #include "TensorOptimization.h"
 #include "release_assert.h"
+#include <format>
 #include <iostream>
 #include <omp.h>
 #include <ranges>
@@ -809,4 +810,29 @@ void mini_jit::TensorOperation::execute_dimension(int64_t index_dim, char const
 mini_jit::TensorConfig mini_jit::TensorOperation::get_config()
 {
   return config;
+}
+
+void mini_jit::TensorOperation::write_kernel_to_file(std::string path_no_extension) const
+{
+  if (prim_first != TensorConfig::prim_t::none && std::holds_alternative<Unary>(first_touch))
+  {
+    std::get<Unary>(first_touch).write_kernel_to_file(std::format("{}_first_touch.bin", path_no_extension).c_str());
+  }
+
+  if (prim_main != TensorConfig::prim_t::none)
+  {
+    if (std::holds_alternative<Brgemm>(main_kernel))
+    {
+      std::get<Brgemm>(main_kernel).write_kernel_to_file(std::format("{}_main.bin", path_no_extension).c_str());
+    }
+    else if (std::holds_alternative<Unary>(main_kernel))
+    {
+      std::get<Unary>(main_kernel).write_kernel_to_file(std::format("{}_main.bin", path_no_extension).c_str());
+    }
+  }
+
+  if (prim_last != TensorConfig::prim_t::none && std::holds_alternative<Unary>(last_touch))
+  {
+    std::get<Unary>(last_touch).write_kernel_to_file(std::format("{}_first_touch.bin", path_no_extension).c_str());
+  }
 }

src/main/TensorOperation.h

@@ -242,6 +242,13 @@ namespace mini_jit
      * @return TensorConfig used by the Tensor operation.
      */
     TensorConfig get_config();
+
+    /**
+     * @brief Writes the current kernel into a file.
+     *
+     * @param path The file to write the kernel to without extension.
+     */
+    void write_kernel_to_file(std::string path_no_extension) const;
   };
 };  // namespace mini_jit
 

src/main/Unary.cpp

@@ -104,4 +104,9 @@ void mini_jit::Unary::relu_unary_fp32(uint32_t m, uint32_t n, uint32_t trans_b)
     kernels::unary_relu(native_kernel, m, n);  // logic of zero_16m_n combined with rest processing
   }
   return;
+}
+
+void mini_jit::Unary::write_kernel_to_file(const char *path) const
+{
+  native_kernel.write(path);
 }

src/main/Unary.h

@@ -92,6 +92,13 @@ class mini_jit::Unary
    * @return pointer to the generated kernel.
    **/
   kernel_t get_kernel() const;
+
+  /**
+   * @brief Writes the current kernel into a file.
+   *
+   * @param path The file to write the kernel to.
+   */
+  void write_kernel_to_file(const char *path) const;
 };
 
 #endif