Add KrnlPrintTensorOp and KrnlPrintOp to facilitate debugging. (llvm#1180)

Signed-off-by: Ettore Tiotto <[email protected]>

Author: Ettore Tiotto

docs/DebuggingNumericalError.md

@@ -82,3 +82,21 @@ optional arguments:
   --rtol RTOL           Relative tolerance for verification
   --atol ATOL           Absolute tolerance for verification
 ```
+
+## Debugging the Code Generated for an Operator.
+
+If you know, or suspect, that a particular ONNX MLIR operator produces an incorrect result, and want to narrow down the problem, we provide a couple of useful Krnl operators that allow printing (at runtime) the value of a tensor, or a value that has a primitive data type. 
+
+To print out the value of a tensor at a particular program point, inject the following code (where `X` is the tensor to be printed):
+
+```code
+create.krnl.printTensor("Tensor X: ", X);
+```
+
+Note: currently the content of the tensor is printed only when the tensor rank is less than four.
+
+To print a message followed by one value, inject the following code (where `val` is the value to be printed and `valType` is its type):
+
+```code
+create.krnl.printf("inputElem: ", val, valType);
+```

include/onnx-mlir/Runtime/OMTensor.h

@@ -284,6 +284,14 @@ int64_t omTensorGetOwning(const OMTensor *tensor);
  */
 void omTensorSetOwning(OMTensor *tensor, int64_t owning);
 
+/**
+ * Print an OMTensor to stdout.
+ *
+ * @param msg, pointer to descriptive string
+ * @param tensor, pointer to the OMTensor to print
+ */
+void omTensorPrint(const char *msg, const OMTensor *tensor);
+
 #ifdef __cplusplus
 }
 #endif

include/onnx-mlir/Runtime/OnnxDataType.h

@@ -26,7 +26,7 @@
 #endif
 
 enum OM_DATA_TYPE {
-#define OM_TYPE_METADATA_DEF(ENUM_NAME, ENUM_VAL, DTYPE_SIZE)                  \
+#define OM_TYPE_METADATA_DEF(ENUM_NAME, ENUM_VAL, DTYPE_SIZE, DTYPE_NAME)      \
   ENUM_NAME = ENUM_VAL,
 #include "OnnxDataTypeMetaData.inc"
 
@@ -38,6 +38,7 @@ typedef enum OM_DATA_TYPE OM_DATA_TYPE;
 #endif
 
 extern const int OM_DATA_TYPE_SIZE[];
+extern const char *OM_DATA_TYPE_NAME[];
 
 #ifdef __cplusplus
 // Note by design const map has no [] operator since [] creates a default
@@ -55,6 +56,9 @@ const std::map<std::string, OM_DATA_TYPE> OM_DATA_TYPE_CPP_TO_ONNX = {
     {"m", ONNX_TYPE_UINT64}, // uint64_t -> UINT64, unsigned long  -> UINT64
     {"f", ONNX_TYPE_FLOAT},  // float    -> FLOAT
     {"d", ONNX_TYPE_DOUBLE}, // double   -> DOUBLE
+    {"PKc", ONNX_TYPE_STRING},    // const char * -> STRING
+    {"Cf", ONNX_TYPE_COMPLEX64},  // _Complex float -> COMPLEX64
+    {"Cd", ONNX_TYPE_COMPLEX128}, // _Complex double -> COMPLEX128
 };
 #endif //__cplusplus
 

include/onnx-mlir/Runtime/OnnxDataTypeMetaData.inc

@@ -3,28 +3,28 @@
  */
 
 #if defined(OM_TYPE_METADATA_DEF)
+// clang-format off
 // Data type metadata declared in the following format:
-// OM_TYPE_METADATA_DEF( dtype enum name, dtype enum value, dtype size)
+// OM_TYPE_METADATA_DEF(dtype enum name, dtype enum value, dtype size, dtype name) 
 // dtype enum values are standard ONNX data types defined in
 // https://github.com/onnx/onnx/blob/main/onnx/onnx.proto#L484
-// clang-format off
-OM_TYPE_METADATA_DEF(ONNX_TYPE_UNDEFINED,   0,  0)
-OM_TYPE_METADATA_DEF(ONNX_TYPE_FLOAT,       1,  sizeof(float))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_UINT8,       2,  sizeof(uint8_t))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_INT8,        3,  sizeof(int8_t))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_UINT16,      4,  sizeof(uint16_t))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_INT16,       5,  sizeof(int16_t))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_INT32,       6,  sizeof(int32_t))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_INT64,       7,  sizeof(int64_t))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_STRING,      8,  0)
-OM_TYPE_METADATA_DEF(ONNX_TYPE_BOOL,        9,  sizeof(bool))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_FLOAT16,     10, 2)
-OM_TYPE_METADATA_DEF(ONNX_TYPE_DOUBLE,      11, sizeof(double))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_UINT32,      12, sizeof(uint32_t))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_UINT64,      13, sizeof(uint64_t))
-OM_TYPE_METADATA_DEF(ONNX_TYPE_COMPLEX64,   14, 8)
-OM_TYPE_METADATA_DEF(ONNX_TYPE_COMPLEX128,  15, 16)
-OM_TYPE_METADATA_DEF(ONNX_TYPE_BFLOAT16,    16, 2)
+OM_TYPE_METADATA_DEF(ONNX_TYPE_UNDEFINED,   0,  0,                "undefined")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_FLOAT,       1,  sizeof(float),    "float")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_UINT8,       2,  sizeof(uint8_t),  "uint8_t")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_INT8,        3,  sizeof(int8_t),   "int8_t")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_UINT16,      4,  sizeof(uint16_t), "uint16_t")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_INT16,       5,  sizeof(int16_t),  "int16_t")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_INT32,       6,  sizeof(int32_t),  "int32_t")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_INT64,       7,  sizeof(int64_t),  "int64_t")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_STRING,      8,  0,                "const char *")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_BOOL,        9,  sizeof(bool),     "_Bool")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_FLOAT16,     10, 2,                "")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_DOUBLE,      11, sizeof(double),   "double")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_UINT32,      12, sizeof(uint32_t), "uint32_t")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_UINT64,      13, sizeof(uint64_t), "uint64_t")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_COMPLEX64,   14, 8,                "_Complex float")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_COMPLEX128,  15, 16,               "_Complex double")
+OM_TYPE_METADATA_DEF(ONNX_TYPE_BFLOAT16,    16, 2,                "")
 // clang-format on
 #else
 #error "Must define OM_TYPE_METADATA_DEF macro."

src/Conversion/KrnlToLLVM/CMakeLists.txt

@@ -2,6 +2,9 @@
 
 add_onnx_mlir_library(OMKrnlToLLVM
   KrnlToLLVM.cpp
+  KrnlToLLVMHelper.cpp  
+  KrnlPrintTensor.cpp
+  KrnlPrint.cpp  
   RuntimeAPI.cpp
 
   LINK_LIBS PUBLIC

src/Conversion/KrnlToLLVM/KrnlPrint.cpp

@@ -0,0 +1,84 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//===------ KrnlPrint.cpp - Lower KrnlPrintOp -----------------------------===//
+//
+// Copyright 2022 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file lowers the KrnlPrintOp operator.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+
+#include "src/Conversion/KrnlToLLVM/KrnlPrint.hpp"
+#include "src/Conversion/KrnlToLLVM/KrnlToLLVM.hpp"
+#include "src/Conversion/KrnlToLLVM/KrnlToLLVMHelper.hpp"
+#include "src/Dialect/Krnl/KrnlHelper.hpp"
+#include "src/Dialect/Krnl/KrnlOps.hpp"
+
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "krnl_to_llvm"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult KrnlPrintOpLowering::matchAndRewrite(Operation *op,
+    ArrayRef<Value> operands, ConversionPatternRewriter &rewriter) const {
+  auto printOp = cast<KrnlPrintOp>(op);
+  Location loc = printOp.getLoc();
+  KrnlPrintOpAdaptor operandAdaptor(operands);
+
+  Value input = operandAdaptor.input();
+  StringRef format = printOp.format();
+  ModuleOp module = printOp->getParentOfType<ModuleOp>();
+
+  // Get a symbol reference to the runtime function to use, creating one if
+  // necessary.
+  auto printfFuncRef = getOrInsertPrintf(rewriter, module);
+
+  // Printf call.
+  LLVM::GlobalOp formatSpec = getOrCreateGlobalString(format, loc, rewriter,
+      module, static_cast<LLVMTypeConverter *>(getTypeConverter()));
+  Value formatSpecPtr = getPtrToGlobalString(formatSpec, loc, rewriter);
+
+  if (input)
+    rewriter.create<CallOp>(loc, printfFuncRef, ArrayRef<Type>({}),
+        ArrayRef<Value>({formatSpecPtr, input}));
+  else
+    rewriter.create<CallOp>(loc, printfFuncRef, ArrayRef<Type>({}),
+        ArrayRef<Value>({formatSpecPtr}));
+
+  rewriter.eraseOp(op);
+  return success();
+}
+
+FlatSymbolRefAttr KrnlPrintOpLowering::getOrInsertPrintf(
+    PatternRewriter &rewriter, ModuleOp module) {
+  // Insert the printf declaration if it is not already present.
+  auto printfFunc = module.lookupSymbol<LLVM::LLVMFuncOp>("printf");
+  MLIRContext *ctx = rewriter.getContext();
+
+  if (!printfFunc) {
+    OpBuilder::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPointToStart(module.getBody());
+    auto voidType = LLVM::LLVMVoidType::get(ctx);
+    Type i8Type = IntegerType::get(ctx, 8);
+    Type i8PtrType = LLVM::LLVMPointerType::get(i8Type);
+    printfFunc =
+        rewriter.create<LLVM::LLVMFuncOp>(rewriter.getUnknownLoc(), "printf",
+            LLVM::LLVMFunctionType::get(voidType, i8PtrType,
+                /*isVarArg=*/true));
+  }
+  return SymbolRefAttr::get(ctx, "printf");
+}
+
+} // namespace onnx_mlir

src/Conversion/KrnlToLLVM/KrnlPrint.hpp

@@ -0,0 +1,44 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//===------ KrnlPrint.hpp - Lower KrnlPrintOp -----------------------------===//
+//
+// Copyright 2022 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file declares the lowering class for the KrnlPrintOp operator.
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Pass/Pass.h"
+
+#include "src/Conversion/KrnlToLLVM/RuntimeAPI.hpp"
+#include "src/Dialect/Krnl/KrnlOps.hpp"
+#include "src/Support/Common.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+class KrnlPrintOpLowering : public ConversionPattern {
+public:
+  explicit KrnlPrintOpLowering(
+      MLIRContext *context, TypeConverter &typeConverter)
+      : ConversionPattern(
+            typeConverter, KrnlPrintOp::getOperationName(), 1, context) {}
+
+  LogicalResult matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const override;
+
+private:
+  static FlatSymbolRefAttr getOrInsertPrintf(
+      PatternRewriter &rewriter, ModuleOp module);
+};
+
+} // namespace onnx_mlir

src/Conversion/KrnlToLLVM/KrnlPrintTensor.cpp

@@ -0,0 +1,75 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//===------ KrnlPrintTensor.cpp - Lower KrnlPrintTensorOp ----------------===//
+//
+// Copyright 2022 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file lowers the KrnlPrintTensorOp operator.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+
+#include "onnx/onnx_pb.h"
+
+#include "src/Conversion/KrnlToLLVM/KrnlPrintTensor.hpp"
+#include "src/Conversion/KrnlToLLVM/KrnlToLLVM.hpp"
+#include "src/Conversion/KrnlToLLVM/KrnlToLLVMHelper.hpp"
+#include "src/Conversion/KrnlToLLVM/RuntimeAPI.hpp"
+#include "src/Dialect/Krnl/KrnlHelper.hpp"
+#include "src/Dialect/Krnl/KrnlOps.hpp"
+
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "krnl_to_llvm"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult KrnlPrintTensorOpLowering::matchAndRewrite(Operation *op,
+    ArrayRef<Value> operands, ConversionPatternRewriter &rewriter) const {
+  auto printTensorOp = cast<KrnlPrintTensorOp>(op);
+  MLIRContext *context = printTensorOp.getContext();
+  Location loc = printTensorOp.getLoc();
+  KrnlPrintTensorOpAdaptor operandAdaptor(operands);
+
+  StringRef msg = printTensorOp.msg();
+  Value input = operandAdaptor.input();
+  assert(input.getType().isa<LLVM::LLVMStructType>() &&
+         "expecting LLVMStructType");
+
+  ModuleOp module = printTensorOp->getParentOfType<ModuleOp>();
+  const auto &apiRegistry = RuntimeAPIRegistry::build(module, rewriter);
+
+  // Get a symbol reference to the runtime function to use, creating one if
+  // necessary.
+  auto int64Ty = IntegerType::get(context, 64);
+  auto memRefTy = input.getType().dyn_cast<LLVM::LLVMStructType>();
+  auto memRefRank = onnx_mlir::getRankFromMemRefType(memRefTy);
+  auto memRefRankVal = rewriter.create<LLVM::ConstantOp>(
+      loc, int64Ty, rewriter.getI64IntegerAttr(memRefRank));
+  Value omTensor = RuntimeAPI::callApi(rewriter, loc, apiRegistry,
+      RuntimeAPI::API::CREATE_OMTENSOR, {memRefRankVal});
+
+  onnx_mlir::fillOMTensorWithMemRef(
+      input, omTensor, false /*outOwning*/, rewriter, loc, apiRegistry, module);
+  LLVM::GlobalOp globalStr = getOrCreateGlobalString(msg, loc, rewriter, module,
+      static_cast<LLVMTypeConverter *>(getTypeConverter()));
+  Value strPtr = getPtrToGlobalString(globalStr, loc, rewriter);
+
+  RuntimeAPI::callApi(rewriter, loc, apiRegistry,
+      RuntimeAPI::API::PRINT_OMTENSOR, {strPtr, omTensor});
+
+  rewriter.eraseOp(op);
+  return success();
+}
+
+} // namespace onnx_mlir

src/Conversion/KrnlToLLVM/KrnlPrintTensor.hpp

@@ -0,0 +1,40 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//===------ KrnlPrintTensor.hpp - Lower KrnlPrintTensorOp -----------------===//
+//
+// Copyright 2022 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file declares the lowering class for the KrnlPrintTensorOp operator.
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Pass/Pass.h"
+
+#include "src/Conversion/KrnlToLLVM/RuntimeAPI.hpp"
+#include "src/Dialect/Krnl/KrnlOps.hpp"
+#include "src/Support/Common.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+class KrnlPrintTensorOpLowering : public ConversionPattern {
+public:
+  explicit KrnlPrintTensorOpLowering(
+      MLIRContext *context, TypeConverter &typeConverter)
+      : ConversionPattern(
+            typeConverter, KrnlPrintTensorOp::getOperationName(), 1, context) {}
+
+  LogicalResult matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const override;
+};
+
+} // namespace onnx_mlir