Dump generated kernels in files

Summary: Dump generated kernels in files for debugging purposes.

Reviewed By: jianyuh

Differential Revision: D14449803

fbshipit-source-id: 58d2b5bc8402ef800a6eeaf573abd2a9ee4f95f4

5 files changed, 91 insertions, 16 deletions

diff --git a/src/GenerateKernel.h b/src/GenerateKernel.h
index 548ed7e..7c0368b 100644
--- a/src/GenerateKernel.h
+++ b/src/GenerateKernel.h
@@ -8,8 +8,10 @@
 #include <asmjit/asmjit.h>
 #include <cpuinfo.h>
 #include <map>
+#include <string>
 #include <tuple>
 #include "fbgemm/Fbgemm.h"
+/*#define FBGEMM_LOG_CODE 1*/
 
 namespace fbgemm {
 
@@ -126,6 +128,32 @@ class CodeGenBase {
       bool accum,
       int leadingDimCRegAssign = 4);
 
+  /**
+   * @brief Generate filename to dump generated code
+   * (debug-only)
+   */
+  template <inst_set_t instSet>
+  std::string getCodeLoggingFile(bool accum, int mc, int nc) {
+    std::string fileName = "gemm_";
+    if (std::is_same<accT, std::int16_t>::value) {
+      fileName += "acc16_";
+    } else if (std::is_same<accT, std::int32_t>::value) {
+      fileName += "acc32_";
+    } else {
+      fileName += "unknown_";
+    }
+    fileName += "accum-" + std::to_string(accum);
+    fileName += "_MC-" + std::to_string(mc);
+    fileName += "_NC-" + std::to_string(nc);
+    if (instSet == inst_set_t::avx512) {
+      fileName += "_avx512";
+    } else if (instSet == inst_set_t::avx2) {
+      fileName += "_avx2";
+    }
+    fileName += ".txt";
+    return fileName;
+  }
+
  private:
   asmjit::X86Ymm
       CRegs_avx2_[12]; ///< AVX2 ymm registers for C in the micro-kernel.
diff --git a/src/GenerateKernelU8S8S32ACC16.cc b/src/GenerateKernelU8S8S32ACC16.cc
index 3419281..e5980b9 100644
--- a/src/GenerateKernelU8S8S32ACC16.cc
+++ b/src/GenerateKernelU8S8S32ACC16.cc
@@ -145,10 +145,16 @@ CodeGenBase<uint8_t, int8_t, int32_t, int16_t>::getOrCreate<inst_set_t::avx2>(
   code_.init(rt_.getCodeInfo());
   asmjit::X86Assembler assembler(&code_);
   asmjit::X86Emitter* a = assembler.asEmitter();
-  // ToDo: Dump in a file for debugging
-  // code dumping/logging
-  // asmjit::FileLogger logger(stderr);
-  // code_.setLogger(&logger);
+
+#if defined(FBGEMM_LOG_CODE)
+  // generated code logging
+  FILE* codeLogfile =
+      fopen(getCodeLoggingFile<inst_set_t::avx2>(accum, mc, nc).c_str(), "w");
+  asmjit::FileLogger* codeLogger = new asmjit::FileLogger(codeLogfile);
+  if (codeLogger) {
+    code_.setLogger(codeLogger);
+  }
+#endif
 
   constexpr int kBlock = PackingTraits<int8_t, int16_t, inst_set_t::avx2>::KCB;
   constexpr int nBlock = PackingTraits<int8_t, int16_t, inst_set_t::avx2>::NCB;
@@ -309,6 +315,12 @@ CodeGenBase<uint8_t, int8_t, int32_t, int16_t>::getOrCreate<inst_set_t::avx2>(
     return nullptr;
   }
   codeCache_[kernelSig] = fn;
+
+#if defined(FBGEMM_LOG_CODE)
+  fclose(codeLogfile);
+  delete codeLogger;
+#endif
+
   return fn;
 }
 
diff --git a/src/GenerateKernelU8S8S32ACC16Avx512.cc b/src/GenerateKernelU8S8S32ACC16Avx512.cc
index ef625e2..cd230c5 100644
--- a/src/GenerateKernelU8S8S32ACC16Avx512.cc
+++ b/src/GenerateKernelU8S8S32ACC16Avx512.cc
@@ -133,10 +133,16 @@ CodeGenBase<uint8_t, int8_t, int32_t, int16_t>::getOrCreate<inst_set_t::avx512>(
   code_.init(rt_.getCodeInfo());
   asmjit::X86Assembler assembler(&code_);
   asmjit::X86Emitter* a = assembler.asEmitter();
-  // ToDo: Dump in a file for debugging
-  // code dumping/logging
-  // asmjit::FileLogger logger(stderr);
-  // code_.setLogger(&logger);
+
+#if defined(FBGEMM_LOG_CODE)
+  // generated code logging
+  FILE* codeLogfile =
+      fopen(getCodeLoggingFile<inst_set_t::avx512>(accum, mc, nc).c_str(), "w");
+  asmjit::FileLogger* codeLogger = new asmjit::FileLogger(codeLogfile);
+  if (codeLogger) {
+    code_.setLogger(codeLogger);
+  }
+#endif
 
   constexpr int kBlock =
       PackingTraits<int8_t, int16_t, inst_set_t::avx512>::KCB;
@@ -301,6 +307,12 @@ CodeGenBase<uint8_t, int8_t, int32_t, int16_t>::getOrCreate<inst_set_t::avx512>(
     return nullptr;
   }
   codeCache_[kernelSig] = fn;
+
+#if defined(FBGEMM_LOG_CODE)
+  fclose(codeLogfile);
+  delete codeLogger;
+#endif
+
   return fn;
 }
 
diff --git a/src/GenerateKernelU8S8S32ACC32.cc b/src/GenerateKernelU8S8S32ACC32.cc
index ab0625c..203dd9a 100644
--- a/src/GenerateKernelU8S8S32ACC32.cc
+++ b/src/GenerateKernelU8S8S32ACC32.cc
@@ -149,10 +149,15 @@ CodeGenBase<uint8_t, int8_t, int32_t, int32_t>::getOrCreate<inst_set_t::avx2>(
   code_.init(rt_.getCodeInfo());
   asmjit::X86Assembler assembler(&code_);
   asmjit::X86Emitter* a = assembler.asEmitter();
-  // ToDo: Dump in a file for debugging
-  // code dumping/logging
-  // asmjit::FileLogger logger(stderr);
-  // code_.setLogger(&logger);
+#if defined(FBGEMM_LOG_CODE)
+  // generated code logging
+  FILE* codeLogfile =
+      fopen(getCodeLoggingFile<inst_set_t::avx2>(accum, mc, nc).c_str(), "w");
+  asmjit::FileLogger* codeLogger = new asmjit::FileLogger(codeLogfile);
+  if (codeLogger) {
+    code_.setLogger(codeLogger);
+  }
+#endif
 
   constexpr int kBlock = PackingTraits<int8_t, int32_t, inst_set_t::avx2>::KCB;
   constexpr int nBlock = PackingTraits<int8_t, int32_t, inst_set_t::avx2>::NCB;
@@ -328,6 +333,12 @@ CodeGenBase<uint8_t, int8_t, int32_t, int32_t>::getOrCreate<inst_set_t::avx2>(
     return nullptr;
   }
   codeCache_[kernelSig] = fn;
+
+#if defined(FBGEMM_LOG_CODE)
+  fclose(codeLogfile);
+  delete codeLogger;
+#endif
+
   return fn;
 }
 
diff --git a/src/GenerateKernelU8S8S32ACC32Avx512.cc b/src/GenerateKernelU8S8S32ACC32Avx512.cc
index e292fa8..0ae0454 100644
--- a/src/GenerateKernelU8S8S32ACC32Avx512.cc
+++ b/src/GenerateKernelU8S8S32ACC32Avx512.cc
@@ -137,10 +137,16 @@ CodeGenBase<uint8_t, int8_t, int32_t, int32_t>::getOrCreate<inst_set_t::avx512>(
   code_.init(rt_.getCodeInfo());
   asmjit::X86Assembler assembler(&code_);
   asmjit::X86Emitter* a = assembler.asEmitter();
-  // ToDo: Dump in a file for debugging
-  // code dumping/logging
-  // asmjit::FileLogger logger(stderr);
-  // code_.setLogger(&logger);
+
+#if defined(FBGEMM_LOG_CODE)
+  // generated code logging
+  FILE* codeLogfile =
+      fopen(getCodeLoggingFile<inst_set_t::avx512>(accum, mc, nc).c_str(), "w");
+  asmjit::FileLogger* codeLogger = new asmjit::FileLogger(codeLogfile);
+  if (codeLogger) {
+    code_.setLogger(codeLogger);
+  }
+#endif
 
   constexpr int kBlock =
       PackingTraits<int8_t, int32_t, inst_set_t::avx512>::KCB;
@@ -319,6 +325,12 @@ CodeGenBase<uint8_t, int8_t, int32_t, int32_t>::getOrCreate<inst_set_t::avx512>(
     return nullptr;
   }
   codeCache_[kernelSig] = fn;
+
+#if defined(FBGEMM_LOG_CODE)
+  fclose(codeLogfile);
+  delete codeLogger;
+#endif
+
   return fn;
 }