[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 ↗

(On Diff #419674)

Please keep these passes sorted.

177 ↗

(On Diff #419674)

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

33

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

213

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

33

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?

213

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();
 
+/// 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.
+void embedMemRefPrints(
+    ModuleOp module,
+    function_ref<SmallVector<Value>(Operation *)> extract_values_func);
+
 //===----------------------------------------------------------------------===//
 // 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/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
+  MLIRLLVMIR
   MLIRLoopLikeInterface
   MLIRMemRef
   MLIRPass
   MLIRTensor
   MLIRTransforms
   MLIRVector
 )
 

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

+//===- EmbedMemRefPrints.cpp - Utility 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/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/IR/Builders.h"
+#include "llvm/Support/FormatVariadic.h"
+
+using namespace mlir;
+
+using LLVM::LLVMFuncOp;
+using ExtractValuesFunc = 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}));
+}
+
+void mlir::memref::embedMemRefPrints(ModuleOp module,
+                                     ExtractValuesFunc extract_values_func) {
+  module.walk([&](FuncOp func) {
+    if (func.isDeclaration())
+      return;
+    Block *body = &func.getBody().front();
+
+    // Print arguments.
+    OpBuilder b(module.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);
+    }
+  });
+}

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/Transforms/embed-memref-prints.mlir

+// RUN: mlir-opt %s --test-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/lib/Dialect/MemRef/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRMemRefTestPasses
+  TestEmbedMemRefPrints.cpp
   TestComposeSubView.cpp
   TestMultiBuffer.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
+  MLIRLLVMIR
   MLIRMemRef
   MLIRMemRefTransforms
   MLIRTestDialect
   )
 
 target_include_directories(MLIRMemRefTestPasses
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../Test
   ${CMAKE_CURRENT_BINARY_DIR}/../Test
   )

mlir/test/lib/Dialect/MemRef/TestEmbedMemRefPrints.cpp

+//===- TestEmbedMemRefPrints.cpp - Test embedded memref prints ------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to test embedded CallOps to print memref content.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+
+using namespace mlir;
+
+static SmallVector<Value> 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()};
+  return {};
+}
+
+namespace {
+
+struct TestEmbedMemRefPrintsPass
+    : public PassWrapper<TestEmbedMemRefPrintsPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestEmbedMemRefPrintsPass)
+
+  StringRef getArgument() const final { return "test-embed-memref-prints"; }
+  StringRef getDescription() const final {
+    return "Test embedding memref prints";
+  }
+
+  void getDependentDialects(DialectRegistry &registry) const {
+    registry.insert<LLVM::LLVMDialect>();
+  }
+
+  void runOnOperation() override {
+    memref::embedMemRefPrints(getOperation(), extractValuesToPrint);
+  }
+};
+
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestEmbedMemRefPrints() {
+  PassRegistration<TestEmbedMemRefPrintsPass>();
+}
+} // namespace test
+} // namespace mlir

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

+// RUN: mlir-opt %s -pass-pipeline="test-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.alloca() : memref<6xf32>
+  %dim = memref.dim %input, %c0 : memref<6xf32>
+  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<6xf32>
+  }
+
+  %output = func.call @reshape(%input)
+  : (memref<6xf32>) -> (memref<2x3xf32>)
+  func.return
+}
+
+
+func.func @reshape(%in: memref<6xf32>) -> memref<2x3xf32> {
+  %out = memref.expand_shape %in [[0, 1]]: memref<6xf32> into memref<2x3xf32>
+  func.return %out : memref<2x3xf32>
+}
+// CHECK: Print memref content after the following op
+// CHECK: @reshape
+// CHECK: rank = 1 offset = 0 sizes = [6] strides = [1]
+// CHECK: [1,  0,  -1,  -2,  -3, -4]
+
+// CHECK: Print memref content after the following op
+// CHECK: memref.expand_shape
+// CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1]
+// CHECK: [1, 0, -1]
+// CHECK: [-2, -3, -4]

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestGpuSerializeToHsacoPass();
 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestComposeSubView();
+void registerTestEmbedMemRefPrints();
 void registerTestMultiBuffering();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgElementwiseFusion();

 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToHsacoPass();
 #endif
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
+  mlir::test::registerTestEmbedMemRefPrints();
   mlir::test::registerTestExpandTanhPass();
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
   mlir::test::registerTestGpuParallelLoopMappingPass();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLinalgCodegenStrategy();

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

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

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

 cc_library(
     name = "TestMemRef",
     srcs = glob(["lib/Dialect/MemRef/*.cpp"]),
     defines = ["MLIR_CUDA_CONVERSIONS_ENABLED"],
     includes = ["lib/Dialect/Test"],
     deps = [
         ":TestDialect",
         "//mlir:Affine",
+        "//mlir:LLVMDialect",
         "//mlir:MemRefDialect",
         "//mlir:MemRefTransforms",
         "//mlir:Pass",
         "//mlir:Transforms",
     ],
 )
 
 cc_library(