Opt reporting: add signatures to stats (#2464)

* Opt reporting: add signatures to stats

Signed-off-by: Alexandre Eichenberger <[email protected]>

---------

Signed-off-by: Alexandre Eichenberger <[email protected]>
Co-authored-by: Tung D. Le <[email protected]>

include/onnx-mlir/Runtime/OMTensor.h

@@ -195,7 +195,8 @@ OM_EXTERNAL_VISIBILITY const int64_t *omTensorGetShape(const OMTensor *tensor);
  *
  * Set the data shape array of the OMTensor to the values in the input array.
  */
-OM_EXTERNAL_VISIBILITY void omTensorSetShape(OMTensor *tensor, const int64_t *shape);
+OM_EXTERNAL_VISIBILITY void omTensorSetShape(
+    OMTensor *tensor, const int64_t *shape);
 
 /**
  * \brief OMTensor data strides getter
@@ -209,7 +210,8 @@ OM_EXTERNAL_VISIBILITY void omTensorSetShape(OMTensor *tensor, const int64_t *sh
  * @param tensor pointer to the OMTensor
  * @return pointer to the data strides array.
  */
-OM_EXTERNAL_VISIBILITY const int64_t *omTensorGetStrides(const OMTensor *tensor);
+OM_EXTERNAL_VISIBILITY const int64_t *omTensorGetStrides(
+    const OMTensor *tensor);
 
 /**
  * \brief OMTensor data strides setter
@@ -320,7 +322,10 @@ OM_EXTERNAL_VISIBILITY void omTensorSetOwning(OMTensor *tensor, int64_t owning);
 /**
  * Print an OMTensor to stdout.
  *
- * @param msg, pointer to descriptive string
+ * @param msg, pointer to descriptive string. It accepts one of 3 formats:
+ * '%t' for printing the tensor's type, '%s' for printing an extensive signature
+ * printout, and '%d' for printing the full data values of the tensor.
+ * Additionally it recognize '%e' as the end of the message string.
  * @param tensor, pointer to the OMTensor to print
  */
 OM_EXTERNAL_VISIBILITY void omTensorPrint(

src/Conversion/ONNXToKrnl/Math/Elementwise.cpp

@@ -2094,7 +2094,7 @@ struct ONNXElementwiseBinaryOpLowering
               shapeHelper, create, op, outputMemRefType,
               collapsedInnermostLoops, estimatedSimdLoopTripCount);
       if (uVL > 0) {
-        if (collapsedInnermostLoops == outputRank)
+        if (collapsedInnermostLoops == (int64_t)outputRank)
           onnxToKrnlSimdReport(op, /*successful*/ true, uVL,
               estimatedSimdLoopTripCount, "binary fully flattened");
         else
@@ -2257,7 +2257,7 @@ struct ONNXElementwiseVariadicOpLowering
               shapeHelper, create, op, outputMemRefType,
               collapsedInnermostLoops, estimatedSimdLoopTripCount);
       if (uVL > 0) {
-        if (collapsedInnermostLoops == outputRank)
+        if (collapsedInnermostLoops == (int64_t)outputRank)
           onnxToKrnlSimdReport(op, /*successful*/ true, uVL,
               estimatedSimdLoopTripCount, "variadic fully flattened");
         else

src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.cpp

@@ -16,7 +16,6 @@
 #include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
-#include "llvm/Support/Path.h"
 
 #include "src/Accelerators/Accelerator.hpp"
 #include "src/Dialect/Krnl/DialectBuilder.hpp"
@@ -664,66 +663,13 @@ bool hasNonIdentityLayout(ValueRange operands) {
 // Support functions for reporting.
 //===----------------------------------------------------------------------===//
 
-// We may try to relate the node names generated by the instrumentation
-// with the node names printed by opt-report. Thus it is key to keep the
-// code that generates these node name in sync.
-//
-// The code are found here:
-// 1) `matchAndRewrite` from `src/Conversion/KrnlToLLVM/KrnlInstrument.cpp`
-// 2) `getNodeNameLikeInKrnlInstrument` from
-//    `src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.cpp`
-
-static std::string getNodeNameLikeInKrnlInstrument(Operation *op) {
-  StringAttr nodeName;
-  // Try with op onnx_node_name attribute.
-  nodeName = op->getAttrOfType<StringAttr>("onnx_node_name");
-  if (nodeName) {
-    LLVM_DEBUG(llvm::dbgs() << "op has node name\n");
-    return nodeName.str();
-  }
-  // Try with op location.
-  Location loc = op->getLoc();
-  if (auto nameLoc = loc.dyn_cast<NameLoc>()) {
-    LLVM_DEBUG(llvm::dbgs() << "op has node name\n");
-    return nameLoc.getName().str();
-  }
-  if (auto fusedLoc = loc.dyn_cast<FusedLoc>()) {
-    // Combine each location name and set it as nodeName.
-    LLVM_DEBUG(llvm::dbgs() << "fuse loc has name\n");
-    std::string name;
-    for (Location locIt : fusedLoc.getLocations()) {
-      if (auto nameLocIt = locIt.dyn_cast<NameLoc>())
-        name += nameLocIt.getName().str() + "-";
-      else if (auto fileLineColLoc = locIt.dyn_cast<FileLineColLoc>()) {
-        std::string filename =
-            llvm::sys::path::filename(fileLineColLoc.getFilename().str()).str();
-        name += filename + ":" + std::to_string(fileLineColLoc.getLine()) + "-";
-      }
-    }
-    if (name.empty())
-      name = "NOTSET";
-    else
-      name.pop_back(); // remove last "-"
-    return name;
-  }
-  if (auto fileLineColLoc = loc.dyn_cast<FileLineColLoc>()) {
-    LLVM_DEBUG(llvm::dbgs() << "file line col has name\n");
-    std::string filename =
-        llvm::sys::path::filename(fileLineColLoc.getFilename().str()).str();
-    std::string name =
-        filename + ":" + std::to_string(fileLineColLoc.getLine());
-    return name;
-  }
-  return "NOTSET";
-}
-
 void impl::onnxToKrnlParallelReport(Operation *op, bool successful,
     int64_t loopLevel, int64_t parallelLoopTripCount,
     const std::string &comment) {
   assert(OnnxToKrnlLoweringConfiguration::reportOnParallel && "must report");
   assert(comment.find(',') == std::string::npos && "no comma in comments");
   StringAttr opName = op->getName().getIdentifier();
-  std::string nodeNameStr = getNodeNameLikeInKrnlInstrument(op);
+  std::string nodeNameStr = getNodeNameInPresenceOfOpt(op);
   // Print report on this op.
   printf("==PAR-REPORT==, %s%s, %s, %s, %lld, %lld\n", opName.data(),
       (successful ? "-parallel" : ""), nodeNameStr.c_str(), comment.c_str(),
@@ -736,7 +682,7 @@ void impl::onnxToKrnlSimdReport(Operation *op, bool successful,
   assert(OnnxToKrnlLoweringConfiguration::reportOnSimd && "must report");
   assert(comment.find(',') == std::string::npos && "no comma in comments");
   StringAttr opName = op->getName().getIdentifier();
-  std::string nodeNameStr = getNodeNameLikeInKrnlInstrument(op);
+  std::string nodeNameStr = getNodeNameInPresenceOfOpt(op);
   // Handling message.
   std::string message = OnnxToKrnlLoweringConfiguration::defaultSimdComment;
   if (message.empty())

src/Conversion/ONNXToKrnl/Tensor/PrintSignature.cpp

@@ -33,15 +33,30 @@ struct ONNXPrintSignatureLowering
     Location loc = ONNXLoc<ONNXPrintSignatureOp>(op);
     MultiDialectBuilder<KrnlBuilder> create(rewriter, loc);
 
+    // First message.
     std::string opName(printSignatureOp.getOpName().data());
-    create.krnl.printf(opName);
-    std::string msg = "%t ";
+    std::string msg =
+        "%i==SIG-REPORT==, " + opName + ", sig"; // meaningless secondary key.
+    // Discover the values to print, setting aside the last one.
+    llvm::SmallVector<Value, 4> printVal;
     for (Value oper : adaptor.getInput())
       if (!oper.getType().isa<NoneType>())
-        create.krnl.printTensor(msg, oper);
-    Value noneValue;
-    rewriter.replaceOpWithNewOp<KrnlPrintOp>(op, "\n", noneValue);
-    // For debug, no need to report on SIMD.
+        printVal.emplace_back(oper);
+    int64_t printNum = printVal.size();
+    Value lastVal = printVal.pop_back_val();
+    // Print all but the last one.
+    for (Value oper : printVal) {
+      create.krnl.printTensor(msg + ", %t%e", oper);
+      msg = "%i";
+    }
+    // Print the last one with replace with new op.
+    if (printNum > 0) {
+      rewriter.replaceOpWithNewOp<KrnlPrintTensorOp>(
+          op, msg + ", %t\n%e", lastVal);
+    } else {
+      Value noneVal = nullptr;
+      rewriter.replaceOpWithNewOp<KrnlPrintTensorOp>(op, msg + "\n%e", noneVal);
+    }
     return success();
   }
 };

src/Dialect/ONNX/ONNXOps/OpHelper.cpp

@@ -14,6 +14,7 @@
 
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Path.h"
 
 #include "src/Dialect/Mlir/IndexExpr.hpp"
 #include "src/Dialect/ONNX/ONNXLayoutHelper.hpp"
@@ -676,4 +677,57 @@ int64_t mlirTypeToOnnxType(Type elemType) {
   return onnxType;
 }
 
+//===----------------------------------------------------------------------===//
+// Support for location.
+//===----------------------------------------------------------------------===//
+
+// We may try to relate the node names generated by the instrumentation
+// with the node names printed by opt-report. Thus it is key to keep the
+// code that generates these node name in sync.
+//
+// The code are found here:
+// 1) `matchAndRewrite` from `src/Conversion/KrnlToLLVM/KrnlInstrument.cpp`
+// 2) `getNodeNameInPresenceOfOpt` from
+//    `src/Dialect/ONNX/ONNXOps/OpHelper.cpp`
+
+std::string getNodeNameInPresenceOfOpt(Operation *op) {
+  StringAttr nodeName;
+  // Try with op onnx_node_name attribute.
+  nodeName = op->getAttrOfType<StringAttr>("onnx_node_name");
+  if (nodeName) {
+    return nodeName.str();
+  }
+  // Try with op location.
+  Location loc = op->getLoc();
+  if (auto nameLoc = loc.dyn_cast<NameLoc>()) {
+    return nameLoc.getName().str();
+  }
+  if (auto fusedLoc = loc.dyn_cast<FusedLoc>()) {
+    // Combine each location name and set it as nodeName.
+    std::string name;
+    for (Location locIt : fusedLoc.getLocations()) {
+      if (auto nameLocIt = locIt.dyn_cast<NameLoc>())
+        name += nameLocIt.getName().str() + "-";
+      else if (auto fileLineColLoc = locIt.dyn_cast<FileLineColLoc>()) {
+        std::string filename =
+            llvm::sys::path::filename(fileLineColLoc.getFilename().str()).str();
+        name += filename + ":" + std::to_string(fileLineColLoc.getLine()) + "-";
+      }
+    }
+    if (name.empty())
+      name = "NOTSET";
+    else
+      name.pop_back(); // remove last "-"
+    return name;
+  }
+  if (auto fileLineColLoc = loc.dyn_cast<FileLineColLoc>()) {
+    std::string filename =
+        llvm::sys::path::filename(fileLineColLoc.getFilename().str()).str();
+    std::string name =
+        filename + ":" + std::to_string(fileLineColLoc.getLine());
+    return name;
+  }
+  return "NOTSET";
+}
+
 } // namespace onnx_mlir

src/Dialect/ONNX/ONNXOps/OpHelper.hpp

@@ -261,4 +261,10 @@ bool areDimsFromConcat(mlir::Value val);
 /// Get all dimensions that are stored by the value.
 void getDims(mlir::Value val, llvm::SmallVectorImpl<mlir::Value> &dims);
 
+//===----------------------------------------------------------------------===//
+// Support for location.
+//===----------------------------------------------------------------------===//
+
+std::string getNodeNameInPresenceOfOpt(mlir::Operation *op);
+
 } // namespace onnx_mlir

src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp

@@ -895,4 +895,5 @@ void SaveOnnxAttrInOp(mlir::Operation *op,
   OP_TYPE specificOp = llvm::cast<OP_TYPE>(op);
   setAttr(specificOp, newAttr);
 }
+
 } // namespace onnx_mlir

src/Runtime/OMInstrument.inc

@@ -80,27 +80,27 @@ static FILE *fout = 0;
 #endif
 
 #ifdef _WIN32
-void TimeInit() {
+static void TimeInit() {
   QueryPerformanceFrequency(&perfFrequency);
   QueryPerformanceCounter(&globalTime);
   initTime = globalTime;
 }
 #else
-void TimeInit() {
+static void TimeInit() {
   gettimeofday(&globalTimeVal, NULL);
   initTimeVal = globalTimeVal;
 }
 #endif
 
 #ifdef _WIN32
-inline void WinTimerSub(LARGE_INTEGER newTime, LARGE_INTEGER prevTime,
+static inline void WinTimerSub(LARGE_INTEGER newTime, LARGE_INTEGER prevTime,
     LONGLONG *resultSeconds, LONGLONG *resultMicroseconds) {
   LONGLONG elapsed = newTime.QuadPart - prevTime.QuadPart;
   *resultSeconds = elapsed / perfFrequency.QuadPart;
   *resultMicroseconds =
       ((elapsed * 1000000) / perfFrequency.QuadPart) % 1000000;
 }
-void ReportTime() {
+static void ReportTime() {
   LARGE_INTEGER newTime;
   LONGLONG resultSeconds, resultMicroseconds;
   QueryPerformanceCounter(&newTime);
@@ -111,7 +111,7 @@ void ReportTime() {
   globalTime = newTime;
 }
 #else
-void ReportTime() {
+static void ReportTime() {
   struct timeval newTimeValue, result;
   gettimeofday(&newTimeValue, NULL);
   timersub(&newTimeValue, &globalTimeVal, &result);
@@ -125,15 +125,15 @@ void ReportTime() {
 #endif
 
 #ifdef _WIN32
-void ReportMemory() {
+static void ReportMemory() {
   PROCESS_MEMORY_COUNTERS_EX pmc;
   GetProcessMemoryInfo(
       GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS *)&pmc, sizeof(pmc));
   SIZE_T vMemSizeKB = pmc.PrivateUsage / 1024;
   fprintf(fout, "%zu\n", vMemSizeKB);
 }
 #else
-void ReportMemory() {
+static void ReportMemory() {
   char memCommand[200];
   char memOutput[200];
   FILE *memPipe;
@@ -159,6 +159,21 @@ void ReportMemory() {
 }
 #endif
 
+FILE *getInstrumentFile() {
+  if (fout)
+    return fout;
+  fout = stdout;
+  if (getenv("ONNX_MLIR_INSTRUMENT_FILE")) {
+    char *fileName = getenv("ONNX_MLIR_INSTRUMENT_FILE");
+    FILE *newFileHandle = fopen(fileName, "w");
+    if (newFileHandle) {
+      fout = newFileHandle;
+    }
+  }
+  assert(fout);
+  return fout;
+}
+
 void OMInstrumentInit() {
   // Read environment variables.
   if (getenv("ONNX_MLIR_NO_INSTRUMENT_TIME")) {
@@ -172,14 +187,7 @@ void OMInstrumentInit() {
     instrumentReportDisabled = true;
   }
   // Handle redirection to file if requested.
-  fout = stdout;
-  if (getenv("ONNX_MLIR_INSTRUMENT_FILE")) {
-    char *fileName = getenv("ONNX_MLIR_INSTRUMENT_FILE");
-    FILE *newFileHandle = fopen(fileName, "w");
-    if (newFileHandle) {
-      fout = newFileHandle;
-    }
-  }
+  getInstrumentFile();
 
   // Init as appropriate.
   if (!instrumentReportDisabled) {