[mlir] Add a pass to print content of memrefs and the related ops IR.

[rriddle]

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

113

This should be camelCase.

121

Drop the llvm here.

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

177

Please keep these passes sorted.

177

This feels like it needs a lot more documentation. It isn't clear from reading this definition what the "printing" is supposed to entail.

mlir/lib/Dialect/MemRef/Transforms/EmbedMemRefPrints.cpp

32

The drop directly to LLVM feels off for this level of abstraction, why is this necessary?

212

This doesn't look great, I don't see why MemRef should be depending on SCF and Linalg here.

[pifon2a]

mlir/lib/Dialect/MemRef/Transforms/EmbedMemRefPrints.cpp

32

This is because I am adding global string constants and printing them. Is there any other way to print a string in MLIR right now?

212

Can't we depend on other dialects in transformations at least?

Just wondering, I removed Linalg and SCF dependency.

[pifon2a]

Address the comments.

[mehdi_amini]

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

122

I'm concerned about C++ injection in passes: it breaks reproducibility.

[pifon2a]

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

122

Fair enough.

It can be fixed by:

• introducing an interface for the ops. I think it is an overkill for this debuggging pass.
• creating a utility function,

mlir::memref::embedMemRefPrints(ExtractValuesFunc extract_values_func)

and creating a test pass that exercises it with the default extract_values_func. After that every user, who wants to print memref content will have to create a pass and call embedMemRefPrints inside runOnOperation with the filter they want. WDYT?

[mehdi_amini]

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

122

utility function + test pass it the common pattern AFAIK

[pifon2a]

Address the concerns. :)

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

 std::unique_ptr<Pass> createResolveRankedShapeTypeResultDimsPass();
 
 /// Creates an operation pass to resolve `memref.dim` operations with values
 /// that are defined by operations that implement the
 /// `InferShapedTypeOpInterface` or the `ReifyRankedShapeTypeShapeOpInterface`,
 /// in terms of shapes of its input operands.
 std::unique_ptr<Pass> createResolveShapedTypeResultDimsPass();
 
+/// Returns a list of memrefs mutated by the given operation.
+SmallVector<Value> ExtractValuesToPrint(Operation *op);
+
+/// Creates a pass to print content of memrefs for every memref argument of a
+/// function and for memref arguments/results of operations in the function
+/// body. What values will be printed is defined by the `extract_values_func`.
+/// If no `extract_values_func` is provided, then the default function will be
+/// used.
+std::unique_ptr<OperationPass<ModuleOp>> createEmbedMemRefPrintsPass(
+    llvm::function_ref<SmallVector<Value>(Operation *)> extract_values_func =
+        ExtractValuesToPrint);
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
 
 } // namespace memref
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES_H

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

let description = [{
     operands.
   }];
   let constructor = "mlir::memref::createResolveShapedTypeResultDimsPass()";
   let dependentDialects = [
     "AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect"
   ];
 }
 
+def EmbedMemRefPrintsPass : Pass<"embed-memref-prints", "mlir::ModuleOp"> {
+  let summary = "Pass to print content of memrefs";
+  let constructor = "mlir::memref::createEmbedMemRefPrintsPass()";
+  let dependentDialects = ["LLVM::LLVMDialect"];
+}
+
 #endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES
 

mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefTransforms
   ComposeSubView.cpp
+  EmbedMemRefPrints.cpp
   ExpandOps.cpp
   FoldSubViewOps.cpp
   MultiBuffer.cpp
   NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAffineUtils
   MLIRArithmetic
   MLIRFunc
   MLIRInferTypeOpInterface
+  MLIRLinalg
   MLIRLoopLikeInterface
   MLIRMemRef
   MLIRPass
+  MLIRSCF
   MLIRTensor
   MLIRTransforms
   MLIRVector
 )
 

mlir/lib/Dialect/MemRef/Transforms/EmbedMemRefPrints.cpp

+//===- EmbedMemRefPrints.cpp - Pass to print content of memrefs  ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "llvm/Support/FormatVariadic.h"
+
+using namespace mlir;
+
+using LLVM::LLVMFuncOp;
+using ExtractValuesFunc = llvm::function_ref<SmallVector<Value>(Operation *)>;
+
+constexpr StringRef kPrintStringFuncName = "print_c_string";
+
+static FlatSymbolRefAttr getOrInsertLLVMFunction(StringRef func_name,
+                                                 Type func_type, Operation *op,
+                                                 OpBuilder *b) {
+  auto module = op->getParentOfType<ModuleOp>();
+  auto tf_func = module.lookupSymbol<LLVMFuncOp>(func_name);
+  if (!tf_func) {
+    OpBuilder::InsertionGuard guard(*b);
+    b->setInsertionPointToStart(module.getBody());
+    tf_func = b->create<LLVMFuncOp>(b->getUnknownLoc(), func_name, func_type);
+  }
+  return SymbolRefAttr::get(b->getContext(), func_name);
+}
+
+static std::string getGlobalName(StringRef base, StringRef content) {
+  return llvm::formatv("{0}_{1}", base, llvm::hash_value(content));
+}
+
+static Value createOrFindGlobalStringConstant(Location loc,
+                                              StringRef global_name,
+                                              StringRef content, OpBuilder *b) {
+  auto module =
+      b->getInsertionBlock()->getParentOp()->getParentOfType<ModuleOp>();
+  Operation *global_constant = SymbolTable::lookupNearestSymbolFrom(
+      module, b->getStringAttr(global_name));
+  if (global_constant) {
+    Value global_ptr = b->create<LLVM::AddressOfOp>(
+        loc, cast<LLVM::GlobalOp>(global_constant));
+    Value c0 =
+        b->create<LLVM::ConstantOp>(loc, b->getI64Type(), b->getIndexAttr(0));
+    return b->create<LLVM::GEPOp>(
+        loc, LLVM::LLVMPointerType::get(b->getIntegerType(8)), global_ptr,
+        ValueRange{c0, c0});
+  }
+  return LLVM::createGlobalString(loc, *b, global_name, content,
+                                  LLVM::Linkage::Internal);
+}
+
+static Operation *emitMemRefPrint(Location loc, Type element_type, Value arg,
+                                  OpBuilder *b) {
+  StringRef func_name;
+  if (element_type.isF32())
+    func_name = "print_memref_f32";
+  if (element_type.isF64())
+    func_name = "print_memref_f64";
+  if (element_type.isInteger(32))
+    func_name = "print_memref_i32";
+  if (element_type.isInteger(64) || element_type.isIndex())
+    func_name = "print_memref_i64";
+  assert(!func_name.empty() &&
+         "Did not find a print function for the element type");
+
+  auto caller_func =
+      b->getInsertionBlock()->getParent()->getParentOfType<FuncOp>();
+  auto func_name_attr = b->getStringAttr(func_name);
+
+  auto callee_func =
+      SymbolTable::lookupNearestSymbolFrom<FuncOp>(caller_func, func_name_attr);
+  if (!callee_func) {
+    OpBuilder::InsertionGuard insertGuard(*b);
+
+    auto module = caller_func->getParentOfType<ModuleOp>();
+    b->setInsertionPointToStart(module.getBody());
+    auto func_type = FunctionType::get(b->getContext(), arg.getType(),
+                                       /*results=*/llvm::None);
+    callee_func = b->create<FuncOp>(module.getLoc(), func_name, func_type);
+    callee_func.setPrivate();
+  }
+  return b->create<func::CallOp>(loc, callee_func, arg);
+}
+
+static bool isElementTypePrintalble(Type element_type) {
+  return element_type.isF32() || element_type.isF64() ||
+         element_type.isInteger(32) || element_type.isInteger(64) ||
+         element_type.isIndex();
+}
+
+static void emitValuePrint(Location loc, Value value, OpBuilder *b) {
+  auto memref_type = value.getType();
+  if (auto unranked_type = memref_type.dyn_cast<UnrankedMemRefType>()) {
+    Type element_type = unranked_type.getElementType();
+    if (!isElementTypePrintalble(element_type))
+      return;
+
+    emitMemRefPrint(loc, element_type, value, b);
+  }
+  if (auto ranked_type = memref_type.dyn_cast<MemRefType>()) {
+    Type element_type = ranked_type.getElementType();
+    if (!isElementTypePrintalble(element_type))
+      return;
+
+    if (element_type.isIndex()) {
+      element_type = b->getI64Type();
+      ranked_type = MemRefType::get(ranked_type.getShape(), element_type,
+                                    ranked_type.getLayout(),
+                                    ranked_type.getMemorySpace());
+      value = b->create<arith::IndexCastOp>(loc, ranked_type, value);
+    }
+
+    auto unranked_type = UnrankedMemRefType::get(
+        element_type, ranked_type.getMemorySpaceAsInt());
+    Value unranked_memref =
+        b->create<memref::CastOp>(loc, unranked_type, value);
+    emitMemRefPrint(loc, element_type, unranked_memref, b);
+  }
+}
+
+static void emitOperationPrint(Operation *op, OpBuilder *b) {
+  std::string debug_str = "\n\nPrint memref content after the following op\n";
+  llvm::raw_string_ostream output_stream(debug_str);
+
+  mlir::OpPrintingFlags flags;
+  op->print(output_stream, flags);
+  output_stream << "\n\n";
+
+  Location loc = op->getLoc();
+  Value message_constant = createOrFindGlobalStringConstant(
+      loc, getGlobalName("debug_op", debug_str), debug_str, b);
+
+  // Insert function call.
+  MLIRContext *ctx = op->getContext();
+  auto func_type = LLVM::LLVMFunctionType::get(
+      LLVM::LLVMVoidType::get(op->getContext()),
+      {LLVM::LLVMPointerType::get(IntegerType::get(ctx, 8))});
+  FlatSymbolRefAttr tf_func_ref =
+      getOrInsertLLVMFunction(kPrintStringFuncName, func_type, op, b);
+  b->create<LLVM::CallOp>(loc, llvm::None, tf_func_ref,
+                          llvm::makeArrayRef({message_constant}));
+}
+
+namespace {
+
+// The pass inserts CallOps to print IR and the content of memrefs for every
+// buffer argument of FuncOp and for every operation within the function
+// depending on the `extract_values_func` filter.
+struct EmbedMemRefPrintsPass
+    : public EmbedMemRefPrintsPassBase<EmbedMemRefPrintsPass> {
+  EmbedMemRefPrintsPass(ExtractValuesFunc extract_values_func)
+      : extract_values_func(extract_values_func) {}
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    module.walk([&](FuncOp func) {
+      if (func.isDeclaration())
+        return;
+      Block *body = &func.getBody().front();
+
+      // Print arguments.
+      OpBuilder b(&getContext());
+      b.setInsertionPointToStart(body);
+      Location loc = func.getLoc();
+      auto args = func.getArguments();
+      if (!args.empty())
+        emitOperationPrint(func, &b);
+      for (auto arg : args)
+        emitValuePrint(loc, arg, &b);
+
+      // Print values extracted from every op.
+      if (!extract_values_func)
+        return;
+      for (auto &op : func.getBody().front().getOperations()) {
+        b.setInsertionPointAfter(&op);
+        auto values = extract_values_func(&op);
+        if (!values.empty())
+          emitOperationPrint(&op, &b);
+        for (Value value : values)
+          emitValuePrint(op.getLoc(), value, &b);
+      }
+    });
+  }
+
+private:
+  ExtractValuesFunc extract_values_func;
+};
+
+} // namespace
+
+SmallVector<Value> mlir::memref::ExtractValuesToPrint(Operation *op) {
+  if (isa<memref::ReinterpretCastOp>(op) || isa<memref::ReshapeOp>(op) ||
+      isa<memref::ExpandShapeOp>(op) || isa<memref::CollapseShapeOp>(op))
+    return {op->getResult(0)};
+  if (auto copy = dyn_cast<memref::CopyOp>(op))
+    return {copy.target()};
+  if (auto loop = dyn_cast<scf::ForOp>(op))
+    return loop.getIterOperands();
+  if (auto linalg = dyn_cast<linalg::LinalgOp>(op))
+    return linalg.getOutputBufferOperands();
+  return {};
+}
+
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::memref::createEmbedMemRefPrintsPass(
+    ExtractValuesFunc extract_values_func) {
+  return std::make_unique<EmbedMemRefPrintsPass>(extract_values_func);
+}

mlir/lib/Dialect/MemRef/Transforms/PassDetail.h

 namespace arith {
 class ArithmeticDialect;
 } // namespace arith
 
 namespace func {
 class FuncDialect;
 } // namespace func
 
+namespace LLVM {
+class LLVMDialect;
+} // namespace LLVM
+  //
 namespace memref {
 class MemRefDialect;
 } // namespace memref
 
 namespace tensor {
 class TensorDialect;
 } // namespace tensor
 

mlir/test/Dialect/MemRef/embed-memref-prints.mlir

+// RUN: mlir-opt %s --embed-memref-prints | FileCheck %s
+
+func.func @print_memrefs(%input: memref<*xf32>) -> memref<*xf32>  {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %rank = memref.rank %input : memref<*xf32>
+  %shape = memref.alloca(%rank) : memref<?xindex>
+  scf.for %i = %c0 to %rank step %c1 {
+    %dim = memref.dim %input, %i : memref<*xf32>
+    memref.store %dim, %shape[%i] : memref<?xindex>
+  }
+
+  %c9000 = arith.constant 9000 : index
+  %num_elem = memref.alloca() : memref<1xindex>
+  memref.store %c9000, %num_elem[%c0] : memref<1xindex>
+  %flat_input = memref.reshape %input(%num_elem)
+    : (memref<*xf32>, memref<1xindex>) -> memref<?xf32>
+
+  %flat_output = memref.alloc(%c9000) : memref<?xf32>
+  %output = memref.reshape %flat_output(%shape)
+    : (memref<?xf32>, memref<?xindex>) -> memref<*xf32>
+  func.return %output : memref<*xf32>
+}
+
+// CHECK-DAG: global internal constant @[[STR0:debug_op_[0-9]+]]({{.*}} @print_memrefs
+// CHECK-DAG: global internal constant @[[STR1:debug_op_[0-9]+]]({{.*}} -> memref<?xf32>
+// CHECK-DAG: global internal constant @[[STR2:debug_op_[0-9]+]]({{.*}} -> memref<*xf32>
+// CHECK-DAG: func private @print_memref_f32(memref<*xf32>)
+// CHECK-DAG: llvm.func @print_c_string(!llvm.ptr<i8>)
+
+// CHECK: func @print_memrefs
+// CHECK-SAME:   (%[[ARG:.*]]: memref<*xf32>)
+// Print debug info for the function arg.
+// CHECK:       %[[STR0_ADDR:.*]] = llvm.mlir.addressof @[[STR0]]
+// CHECK:       %[[STR0_PTR:.*]] = llvm.getelementptr %[[STR0_ADDR]]
+// CHECK:       llvm.call @print_c_string(%[[STR0_PTR]]) : (!llvm.ptr<i8>)
+// CHECK:       call @print_memref_f32(%[[ARG]]) : (memref<*xf32>) -> ()
+
+// Print debug info for reshape from unranked to ranked.
+// CHECK:       %[[RESHAPE:.*]] = memref.reshape %[[ARG]]
+// CHECK:       %[[STR1_ADDR:.*]] = llvm.mlir.addressof @[[STR1]]
+// CHECK:       %[[STR1_PTR:.*]] = llvm.getelementptr %[[STR1_ADDR]]
+// CHECK:       llvm.call @print_c_string(%[[STR1_PTR]]) : (!llvm.ptr<i8>)
+// CHECK:       %[[UNRANKED_BUF:.*]] = memref.cast %[[RESHAPE]]
+// CHECK:       call @print_memref_f32(%[[UNRANKED_BUF]]) : (memref<*xf32>)
+
+// Print debug info for reshape from ranked to unranked.
+// CHECK:       %[[ALLOC:.*]] = memref.alloc
+// CHECK:       %[[RESHAPE_2:.*]] = memref.reshape %[[ALLOC]]
+// CHECK:       %[[STR2_ADDR:.*]] = llvm.mlir.addressof @[[STR2]]
+// CHECK:       %[[STR2_PTR:.*]] = llvm.getelementptr %[[STR2_ADDR]]
+// CHECK:       llvm.call @print_c_string(%[[STR2_PTR]]) : (!llvm.ptr<i8>)
+// CHECK:       call @print_memref_f32(%[[RESHAPE_2]]) : (memref<*xf32>)

mlir/test/mlir-cpu-runner/debug_memref_content.mlir

+// RUN: mlir-opt %s -pass-pipeline="embed-memref-prints,func.func(convert-linalg-to-loops,convert-scf-to-cf),convert-linalg-to-llvm,convert-math-to-llvm,convert-memref-to-llvm,func.func(convert-arith-to-llvm),convert-func-to-llvm,reconcile-unrealized-casts" |\
+// RUN: mlir-cpu-runner -e main -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext |\
+// RUN: FileCheck %s
+
+#map = affine_map<(d0) -> (d0)>
+func.func @main() -> () {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+
+  // Initialize input.
+  %input = memref.alloc() : memref<5xf32>
+  %dim = memref.dim %input, %c0 : memref<5xf32>
+  scf.for %i = %c0 to %dim step %c1 {
+    %diff = arith.subi %c1, %i : index
+    %diff_i64 = arith.index_cast %diff : index to i64
+    %diff_f32 = arith.sitofp %diff_i64 : i64 to f32
+    memref.store %diff_f32, %input[%i] : memref<5xf32>
+  }
+
+  %output = func.call @compute_abs(%input) : (memref<5xf32>) -> (memref<5xf32>)
+  memref.dealloc %output : memref<5xf32>
+  func.return
+}
+
+
+func.func @compute_abs(%in: memref<5xf32>) -> memref<5xf32> {
+  %out = memref.alloc() : memref<5xf32>
+  linalg.generic { indexing_maps = [#map, #map],
+                   iterator_types = ["parallel"]}
+        ins(%in : memref<5xf32>) outs(%out : memref<5xf32>) {
+  ^bb0(%arg2: f32, %arg3: f32):
+    %1 = math.abs %arg2 : f32
+    linalg.yield %1 : f32
+  }
+  func.return %out : memref<5xf32>
+}
+// CHECK: Print memref content after the following op
+// CHECK: @compute_abs
+// CHECK: rank = 1 offset = 0 sizes = [5] strides = [1]
+// CHECK: [1,  0,  -1,  -2,  -3]
+
+// CHECK: Print memref content after the following op
+// CHECK: linalg.generic
+// CHECK:   math.abs
+// CHECK: rank = 1 offset = 0 sizes = [5] strides = [1]
+// CHECK: [1,  0,  1,  2,  3]

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     deps = [
         ":Affine",
         ":AffineUtils",
         ":ArithmeticDialect",
         ":ArithmeticTransforms",
         ":FuncDialect",
         ":IR",
         ":InferTypeOpInterface",
+        ":LLVMDialect",
+        ":LinalgOps",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":MemRefPassIncGen",
         ":Pass",
+        ":SCFDialect",
         ":TensorDialect",
         ":Transforms",
         ":VectorOps",
         "//llvm:Support",
     ],
 )
 
 td_library(