Extend instrumentSignature to print data

src/Accelerators/NNPA/Compiler/NNPACompilerUtils.cpp

@@ -260,8 +260,7 @@ void addPassesNNPA(mlir::OwningOpRef<mlir::ModuleOp> &module,
       else if (optStr == "-O3")
         optLevel = OptLevel::O3;
       // Lower ONNX to Krnl, ZHigh to ZLow.
-      addONNXToKrnlPasses(
-          pm, optLevel, /*enableCSE*/ true, instrumentSignatures, ONNXOpStats);
+      addONNXToKrnlPasses(pm, optLevel, /*enableCSE*/ true, ONNXOpStats);
 
       if (nnpaEmissionTarget >= EmitZLowIR)
         emissionTarget = EmitMLIR;

src/Compiler/CompilerOptions.cpp

@@ -501,14 +501,18 @@ static llvm::cl::opt<std::string, true> parallelizeOpsOpt("parallelize-ops",
 
 static llvm::cl::opt<std::string, true> instrumentSignatureOpt(
     "instrument-signature",
-    llvm::cl::desc("Specify which high-level operations should print their"
-                   " input type(s) and shape(s)\n"
-                   "\"ALL\" or \"\" for all available operations.\n"
-                   "\"NONE\" for no instrument (default).\n"
-                   "\"ops1,ops2, ...\" for the multiple ops.\n"
-                   "e.g. \"onnx.MatMul,onnx.Add\" for MatMul and Add ops.\n"
-                   "Asterisk is also available.\n"
-                   "e.g. \"onnx.*\" for all onnx operations.\n"),
+    llvm::cl::desc(
+        "Specify which high-level operations should print their"
+        " input type(s) and shape(s)\n"
+        "\"ALL\" or \"\" for all available operations.\n"
+        "\"NONE\" for no instrument (default).\n"
+        "\"ops1,ops2, ...\" for the multiple ops.\n"
+        "e.g. \"onnx.MatMul,onnx.Add\" for MatMul and Add ops.\n"
+        "Asterisk is also available.\n"
+        "e.g. \"onnx.*\" for all onnx operations.\n"
+        "If this option is started with \"onnx_node_name\"\n"
+        "the attribute of \"onnx_node_name\", instead of the op name\n"
+        "will be used to match the op, and the data value will be printed.\n"),
     llvm::cl::location(instrumentSignatures), llvm::cl::init("NONE"),
     llvm::cl::cat(OnnxMlirOptions));
 

src/Compiler/CompilerPasses.cpp

@@ -169,10 +169,16 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
   if (instrumentStage == onnx_mlir::InstrumentStages::Onnx)
     pm.addNestedPass<func::FuncOp>(
         onnx_mlir::createInstrumentPass(instrumentOps, instrumentActions));
+  // Print Signatures of each op at runtime if enabled. Should not run
+  // signature and instrument passes at the same time as time may include printf
+  // overheads.
+  if (instrumentSignatures != "NONE")
+    pm.addNestedPass<func::FuncOp>(
+        onnx_mlir::createInstrumentONNXSignaturePass(instrumentSignatures));
 }
 
 void addONNXToKrnlPasses(mlir::PassManager &pm, int optLevel, bool enableCSE,
-    std::string instrumentSignatureString, std::string ONNXOpsStatFormat) {
+    std::string ONNXOpsStatFormat) {
   if (enableCSE)
     // Eliminate common sub-expressions before lowering to Krnl.
     // TODO: enable this by default when we make sure it works flawlessly.
@@ -196,12 +202,6 @@ void addONNXToKrnlPasses(mlir::PassManager &pm, int optLevel, bool enableCSE,
     }
   }
 
-  // Print Signatures of each op at runtime if enabled. Should not run
-  // signature and instrument passes at the same time as time may include printf
-  // overheads.
-  if (instrumentSignatureString != "NONE")
-    pm.addNestedPass<func::FuncOp>(onnx_mlir::createInstrumentONNXSignaturePass(
-        instrumentSignatureString));
   pm.addPass(onnx_mlir::createLowerToKrnlPass(/*enableTiling*/ optLevel >= 3,
       /*enableSIMD*/ optLevel >= 3 && !disableSimdOption, enableParallel,
       /*enableFastMath*/ optLevel >= 3 && enableFastMathOption,
@@ -214,7 +214,7 @@ void addONNXToKrnlPasses(mlir::PassManager &pm, int optLevel, bool enableCSE,
 
 void addKrnlToAffinePasses(mlir::PassManager &pm) {
   pm.addNestedPass<func::FuncOp>(
-      onnx_mlir::krnl::createConvertKrnlToAffinePass(enableParallel));
+      onnx_mlir::krnl::createConvertKrnlToAffinePass());
 }
 
 void addKrnlToLLVMPasses(
@@ -325,8 +325,8 @@ void addPasses(mlir::OwningOpRef<ModuleOp> &module, mlir::PassManager &pm,
 
   if (emissionTarget >= EmitMLIR) {
     if (inputIRLevel <= ONNXLevel)
-      addONNXToKrnlPasses(pm, OptimizationLevel, /*enableCSE*/ true,
-          instrumentSignatures, ONNXOpStats);
+      addONNXToKrnlPasses(
+          pm, OptimizationLevel, /*enableCSE*/ true, ONNXOpStats);
     if (inputIRLevel <= MLIRLevel)
       addKrnlToAffinePasses(pm);
   }

src/Compiler/CompilerPasses.hpp

@@ -23,7 +23,7 @@ void configurePasses();
 void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
     bool donotScrubDisposableElementsAttr = false);
 void addONNXToKrnlPasses(mlir::PassManager &pm, int optLevel, bool enableCSE,
-    std::string instrumentSignatureString, std::string ONNXOpsStatFilename);
+    std::string ONNXOpsStatFilename);
 void addKrnlToAffinePasses(mlir::PassManager &pm);
 void addKrnlToLLVMPasses(
     mlir::OpPassManager &pm, std::string outputNameNoExt, bool enableCSE);

src/Conversion/ONNXToKrnl/Tensor/PrintSignature.cpp

@@ -50,15 +50,22 @@ struct ONNXPrintSignatureLowering
           op, msg + "(no tensors)\n%e", noneVal);
       return success();
     }
+    // Control how the tensor will be printed
+    // Print the only the shape.
+    std::string printControl = "%e";
+    if (printSignatureOp.getPrintData() == 1) {
+      // The data of tensor will be printed
+      printControl = "%d";
+    }
     Value lastVal = printVal.pop_back_val();
     // Print all but the last one.
     for (Value oper : printVal) {
-      create.krnl.printTensor(msg + ", %t%e", oper);
+      create.krnl.printTensor(msg + ", %t" + printControl, oper);
       msg = "%i";
     }
     // Print the last one with replace with new op.
     rewriter.replaceOpWithNewOp<KrnlPrintTensorOp>(
-        op, msg + ", %t\n%e", lastVal);
+        op, msg + ", %t\n" + printControl, lastVal);
     return success();
   }
 };

src/Dialect/ONNX/AdditionalONNXOps.td

@@ -322,15 +322,19 @@ def ONNXNoneOp : ONNX_Op<"NoValue", [ConstantLike, Pure]> {
 //===----------------------------------------------------------------------===//
 // ONNX PrintSignatureOp.
 def ONNXPrintSignatureOp:ONNX_Op<"PrintSignature", []> {
-  let summary = "ONNX Op to print type signature of its input operands";
+  let summary = "ONNX Op to print type signature or data of its input operands";
   let description = [{
-    Print type signature of the op's input operands. This operation is introduced early
-    so as to preserve the name of the original ONNX op.
+    Print type signature of the op's input operands.
+    The parameter op_name specifies a string to be printed with the output, and usually the op_name and onnx_node_name are used.
+    This operation is introduced early so as to preserve the name of the original ONNX op.
+    The argument print_data control whether signature of data of the tensors to be printed.
+    When print_data == 1, the data of the tensor will be printed. Otherwise, the shape of the tensor will be printed.
+    The argument input specifies all the tensor to be printed.
 
     This operation is not part of the standard and was added to assist onnx-mlir.
   }];
 
-  let arguments = (ins StrAttr:$op_name, Variadic<AnyTypeOf<[AnyTensor, NoneType]>>:$input);
+  let arguments = (ins StrAttr:$op_name, SI64Attr:$print_data, Variadic<AnyTypeOf<[AnyTensor, NoneType]>>:$input);
 }
 
 //===----------------------------------------------------------------------===//

src/Dialect/ONNX/Transforms/InstrumentONNXSignaturePass.cpp

@@ -72,15 +72,34 @@ class InstrumentONNXSignaturePass
   void runOnOperation() override {
     onnx_mlir::EnableByRegexOption traceSpecificOpPattern(
         /*emptyIsNone*/ false);
+    // If the signaturePattern is started with "onnx_node_name: ",
+    // the node will be selected by the string attribute of onnx_node_name.
+    std::string header = "onnx_node_name:";
+    std::cout << signaturePattern << "\n";
+    bool useNodeName = false;
+    if (signaturePattern.rfind(header, 0) == 0) {
+      signaturePattern.erase(0, header.length());
+      useNodeName = true;
+      std::cout << signaturePattern << "\n";
+    }
+
     traceSpecificOpPattern.setRegexString(signaturePattern);
     // Iterate on the operations nested in this function.
     getOperation().walk([&](mlir::Operation *op) {
       std::string opName = op->getName().getStringRef().str();
+      std::string matchString = opName;
+      if (useNodeName) {
+        StringAttr onnxNodeName =
+            op->getAttrOfType<mlir::StringAttr>("onnx_node_name");
+        if (onnxNodeName && !onnxNodeName.getValue().empty()) {
+          matchString = onnxNodeName.getValue().str();
+        }
+      }
       auto dialect = op->getDialect();
       if (isa<func::FuncDialect>(dialect) || isa<ONNXPrintSignatureOp>(op)) {
         // Always skip function dialects (such as function call/return), as well
         // as ONNX print signature ops.
-      } else if (traceSpecificOpPattern.isEnabled(opName)) {
+      } else if (traceSpecificOpPattern.isEnabled(matchString)) {
         // Add signature printing op.
         Location loc = op->getLoc();
         OpBuilder builder(op);
@@ -94,7 +113,9 @@ class InstrumentONNXSignaturePass
         // Since we may use the result of an operation, we must insert the
         // print operation after the operation.
         builder.setInsertionPointAfter(op);
-        builder.create<ONNXPrintSignatureOp>(loc, fullNameAttr, operAndRes);
+        // When one node is selected, print the details of the tensor.
+        builder.create<ONNXPrintSignatureOp>(
+            loc, fullNameAttr, useNodeName ? 1 : 0, operAndRes);
       }
     });
   }