[mlir][linalg][bufferize] Debug output as IR attributes

[springerm]

Summary

Instead of printing analysis debug information to stderr, annotate the IR. This makes it easier to understand decisions made by the analysis, especially in larger input IR.

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

struct BufferizationOptions {
   /// Seed for the analysis fuzzer. If set to `0`, the fuzzer is deactivated.
   /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
 
   /// If set to `true`, does not modify the IR apart from adding attributes (for
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
+  /// If set to `true`, the IR is annotated with details about RaW conflicts.
+  /// For debugging only. Should be used together with `testAnalysisOnly`.
+  bool printConflicts = false;
+
   /// Registered post analysis steps.
   PostAnalysisStepList postAnalysisSteps;
 
   /// Only bufferize ops from dialects that are allowed-listed by the filter.
   /// All other ops are ignored. This option controls the scope of partial
   /// bufferization.
   ///
   /// Note: If no filter is specified, all ops are bufferized (as long as they

mlir/include/mlir/Dialect/Linalg/Passes.td

let description = [{
     function boundaries. Within a function boundary, the analysis is performed
     on SSA use-def chains starting from function operands that are annotated
     with the 'inplaceable' attribute.
   }];
   let options = [
     Option<"testAnalysisOnly", "test-analysis-only", "bool",
             /*default=*/"false",
            "Only runs inplaceability analysis (for testing purposes only)">,
+    Option<"printConflicts", "print-conflicts", "bool",
+            /*default=*/"false",
+           "Annotates IR with RaW conflicts. Requires test-analysis-only.">,
     Option<"allowReturnMemref", "allow-return-memref", "bool",
             /*default=*/"false",
            "Allows the return of memrefs (for testing purposes only)">,
     Option<"allowUnknownOps", "allow-unknown-ops", "bool",
            /*default=*/"false",
            "Allows unknown (not bufferizable) ops in the input IR.">,
     Option<"useAlloca", "use-alloca", "bool",
            /*default=*/"false",

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/FormatVariadic.h"
-
-#define DEBUG_TYPE "comprehensive-module-bufferize"
 
 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 using namespace comprehensive_bufferize;
 
-#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
-#define LDBG(X) LLVM_DEBUG(DBGS() << X)
-
-// Forward declarations.
-#ifndef NDEBUG
-static std::string printOperationInfo(Operation *, bool prefix = true);
-static std::string printValueInfo(Value, bool prefix = true);
-#endif
-
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific attribute manipulation.
 // These are for testing and debugging only. Bufferization information is
 // stored in BufferizationAliasInfo. When run with `testAnalysisOnly`, the IR
 // is annotated with the results of the analysis (copied from
 // BufferizationAliasInfo), so that they can be checked in tests.
static void setInPlaceOpResult(OpResult opResult, bool inPlace) {
 
   Operation *op = opResult.getOwner();
   auto attr =
       op->getAttr(kInPlaceResultsAttrName).dyn_cast_or_null<ArrayAttr>();
   SmallVector<StringRef> inPlaceVector =
       attr ? SmallVector<StringRef>(
                  llvm::to_vector<4>(attr.getAsValueRange<StringAttr>()))
            : SmallVector<StringRef>(op->getNumResults(), "false");
-  LDBG("->set inPlace=" << inPlace << " <- #" << opResult.getResultNumber()
-                        << ": " << printOperationInfo(op) << "\n");
   inPlaceVector[opResult.getResultNumber()] = inPlace ? "true" : "false";
   op->setAttr(kInPlaceResultsAttrName,
               OpBuilder(op).getStrArrayAttr(inPlaceVector));
 }
 
 //===----------------------------------------------------------------------===//
-// Printing helpers.
-//===----------------------------------------------------------------------===//
-
-#ifndef NDEBUG
-/// Helper method printing the bufferization information of a buffer / tensor.
-static void printTensorOrBufferInfo(std::string prefix, Value value,
-                                    AsmState &state, llvm::raw_ostream &os) {
-  if (!value.getType().isa<ShapedType>())
-    return;
-  os << prefix;
-  value.printAsOperand(os, state);
-  os << " : " << value.getType();
-}
-
-/// Print the operation name and bufferization information.
-static std::string printOperationInfo(Operation *op, bool prefix) {
-  std::string result;
-  llvm::raw_string_ostream os(result);
-  AsmState state(op->getParentOfType<mlir::FuncOp>());
-  StringRef tab = prefix ? "\n[" DEBUG_TYPE "]\t" : "";
-  os << tab << op->getName();
-  SmallVector<Value> shapedOperands;
-  for (OpOperand &opOperand : op->getOpOperands()) {
-    std::string prefix =
-        llvm::formatv("{0}  -> #{1} ", tab, opOperand.getOperandNumber());
-    printTensorOrBufferInfo(prefix, opOperand.get(), state, os);
-  }
-  for (OpResult opResult : op->getOpResults()) {
-    std::string prefix =
-        llvm::formatv("{0}  <- #{1} ", tab, opResult.getResultNumber());
-    printTensorOrBufferInfo(prefix, opResult, state, os);
-  }
-  return result;
-}
-
-/// Print the bufferization information for the defining op or block argument.
-static std::string printValueInfo(Value value, bool prefix) {
-  auto *op = value.getDefiningOp();
-  if (op)
-    return printOperationInfo(op, prefix);
-  // Print the block argument bufferization information.
-  std::string result;
-  llvm::raw_string_ostream os(result);
-  AsmState state(value.getParentRegion()->getParentOfType<mlir::FuncOp>());
-  os << value;
-  printTensorOrBufferInfo("\n\t - ", value, state, os);
-  return result;
-}
-#endif
-
-//===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
 //===----------------------------------------------------------------------===//
 
 /// Return true if opOperand has been decided to bufferize in-place.
 static bool isInplaceMemoryWrite(OpOperand &opOperand,
                                  const BufferizationAliasInfo &aliasInfo) {
   // The analysis does not know what happens to the result of a ToMemrefOp, so
   // we assume that it is written to.
static bool isInplaceMemoryWrite(OpOperand &opOperand,
   return aliasInfo.isInPlace(opResult);
 }
 
 /// Return true if, under current bufferization decisions, the buffer of `value`
 /// is not writable.
 static bool aliasesNonWritableBuffer(Value value,
                                      const BufferizationAliasInfo &aliasInfo,
                                      BufferizationState &state) {
-  LDBG("WRITABILITY ANALYSIS FOR " << printValueInfo(value) << "\n");
   bool foundNonWritableBuffer = false;
   aliasInfo.applyOnAliases(value, [&](Value v) {
     // Query BufferizableOpInterface to see if the OpResult is writable.
     // TODO: Out-of-place bufferized OpResult could be considered writable.
     if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(v))
       if (bufferizableOp && bufferizableOp.isWritable(v, state))
         return;
 
     // Query BufferizableOpInterface to see if the BlockArgument is writable.
     if (auto bbArg = v.dyn_cast<BlockArgument>())
       if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(
               bbArg.getOwner()->getParentOp()))
         if (bufferizableOp.isWritable(bbArg, state))
           return;
 
     foundNonWritableBuffer = true;
   });
 
-  if (foundNonWritableBuffer)
-    LDBG("--> NON WRITABLE\n");
-  else
-    LDBG("--> WRITABLE\n");
-
   return foundNonWritableBuffer;
 }
 
 /// Return true if the buffer to which `operand` would bufferize is equivalent
 /// to some buffer write.
 static bool aliasesInPlaceWrite(Value value,
                                 const BufferizationAliasInfo &aliasInfo) {
-  LDBG("----Start aliasesInPlaceWrite\n");
-  LDBG("-------for : " << printValueInfo(value) << '\n');
   bool foundInplaceWrite = false;
   aliasInfo.applyOnAliases(value, [&](Value v) {
     for (auto &use : v.getUses()) {
       if (isInplaceMemoryWrite(use, aliasInfo)) {
-        LDBG("-----------wants to bufferize to inPlace write: "
-             << printOperationInfo(use.getOwner()) << '\n');
         foundInplaceWrite = true;
         return;
       }
     }
   });
-
-  if (!foundInplaceWrite)
-    LDBG("----------->does not alias an inplace write\n");
-
   return foundInplaceWrite;
 }
 
 /// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors
 /// properly dominates `b` and `b` is not inside `a`.
 static bool happensBefore(Operation *a, Operation *b,
                           const DominanceInfo &domInfo) {
   do {
     // TODO: Instead of isProperAncestor + properlyDominates, we should use
     // properlyDominatesImpl(a, b, /*enclosingOpOk=*/false)
     if (a->isProperAncestor(b))
       return false;
     if (domInfo.properlyDominates(a, b))
       return true;
   } while ((a = a->getParentOp()));
   return false;
 }
 
+/// Annotate IR with details about the detected RaW conflict.
+static void annotateConflict(OpOperand *uRead, OpOperand *uConflictingWrite,
+                             Value lastWrite) {
+  static uint64_t counter = 0;
+  Operation *readingOp = uRead->getOwner();
+  Operation *conflictingWritingOp = uConflictingWrite->getOwner();
+
+  OpBuilder b(conflictingWritingOp->getContext());
+  std::string id = "C_" + std::to_string(counter++);
+
+  std::string conflictingWriteAttr =
+      id +
+      "[CONFL-WRITE: " + std::to_string(uConflictingWrite->getOperandNumber()) +
+      "]";
+  conflictingWritingOp->setAttr(conflictingWriteAttr, b.getUnitAttr());
+
+  std::string readAttr =
+      id + "[READ: " + std::to_string(uRead->getOperandNumber()) + "]";
+  readingOp->setAttr(readAttr, b.getUnitAttr());
+
+  if (auto opResult = lastWrite.dyn_cast<OpResult>()) {
+    std::string lastWriteAttr = id + "[LAST-WRITE: result " +
+                                std::to_string(opResult.getResultNumber()) +
+                                "]";
+    opResult.getDefiningOp()->setAttr(lastWriteAttr, b.getUnitAttr());
+  } else {
+    auto bbArg = lastWrite.cast<BlockArgument>();
+    std::string lastWriteAttr =
+        id + "[LAST-WRITE: bbArg " + std::to_string(bbArg.getArgNumber()) + "]";
+    bbArg.getOwner()->getParentOp()->setAttr(lastWriteAttr, b.getUnitAttr());
+  }
+}
+
 /// Given sets of uses and writes, return true if there is a RaW conflict under
 /// the assumption that all given reads/writes alias the same buffer and that
 /// all given writes bufferize inplace.
 ///
 /// A conflict is: According to SSA use-def chains, a read R is supposed to read
 /// the result of a write W1. But because of bufferization decisions, R actually
 /// reads another write W2.
 static bool hasReadAfterWriteInterference(
for (OpOperand *uRead : usesRead) {
 
     // Look for conflicting memory writes. Potential conflicts are writes to an
     // alias that have been decided to bufferize inplace.
     for (OpOperand *uConflictingWrite : usesWrite) {
       // Throughout this loop, check for multiple requirements that have to be
       // met for uConflictingWrite to be an actual conflict.
       Operation *conflictingWritingOp = uConflictingWrite->getOwner();
 
-      // Print some debug info.
-      LDBG("Found potential conflict:\n");
-      LDBG("READ = #" << uRead->getOperandNumber() << " of "
-                      << printOperationInfo(readingOp) << "\n");
-      LDBG("CONFLICTING WRITE = #"
-           << uConflictingWrite->getOperandNumber() << " of "
-           << printOperationInfo(conflictingWritingOp) << "\n");
-
       // No conflict if the readingOp dominates conflictingWritingOp, i.e., the
       // write is not visible when reading.
       if (happensBefore(readingOp, conflictingWritingOp, domInfo))
         continue;
 
       // No conflict if the reading use equals the use of the conflicting write.
       // A use cannot conflict with itself. Note: Just being the same op is not
       // enough. It has to be the same use.
for (OpOperand *uConflictingWrite : usesWrite) {
           if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state,
                                               aliasInfo))
             continue;
 
       // Ops are not conflicting if they are in mutually exclusive regions.
       if (insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp))
         continue;
 
-      LDBG("WRITE = #" << printValueInfo(lastWrite) << "\n");
-
       // No conflict if the conflicting write happens before the last
       // write.
       if (Operation *writingOp = lastWrite.getDefiningOp()) {
         if (happensBefore(conflictingWritingOp, writingOp, domInfo))
           // conflictingWritingOp happens before writingOp. No conflict.
           continue;
         // No conflict if conflictingWritingOp is contained in writingOp.
         if (writingOp->isProperAncestor(conflictingWritingOp))
           continue;
       } else {
         auto bbArg = lastWrite.cast<BlockArgument>();
         Block *block = bbArg.getOwner();
         if (!block->findAncestorOpInBlock(*conflictingWritingOp))
           // conflictingWritingOp happens outside of the block. No
           // conflict.
           continue;
       }
 
       // No conflict if the conflicting write and the last write are the same
       // use.
       if (getAliasingOpResult(*uConflictingWrite) == lastWrite)
         continue;
 
       // All requirements are met. Conflict found!
-      LDBG("CONFLICT CONFIRMED!\n\n");
+
+      if (options.printConflicts)
+        annotateConflict(uRead, uConflictingWrite, lastWrite);
+
       return true;
     }
   }
 
-  LDBG("NOT A CONFLICT!\n\n");
   return false;
 }
 
 /// Return true if bufferizing result inplace would create a conflict. A read R
 /// and a write W of the same alias set is a conflict if inplace bufferization
 /// of W changes the value read by R to a value different from the one that
 /// would be expected by tracing back R's origin through SSA use-def chains.
 /// A conflict can only be introduced by a new alias and/or an inplace
#endif // NDEBUG
 
   // This is a problem only if the buffer is written to via some alias.
   bool hasWrite = aliasesInPlaceWrite(opResult, aliasInfo) ||
                   aliasesInPlaceWrite(opOperand.get(), aliasInfo) ||
                   bufferizesToMemoryWrite(opOperand);
   if (!hasWrite)
     return false;
 
-  LDBG("->the corresponding buffer is not writeable\n");
   return true;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 /// Determine if `operand` can be bufferized in-place with `result`.
 static LogicalResult bufferizableInPlaceAnalysisImpl(
     OpOperand &operand, OpResult result, BufferizationAliasInfo &aliasInfo,
     BufferizationState &state, const DominanceInfo &domInfo) {
 #ifndef NDEBUG
   SmallVector<OpOperand *> opOperands = getAliasingOpOperand(result);
   assert(llvm::find(opOperands, &operand) != opOperands.end() &&
          "operand and result do not match");
 #endif // NDEBUG
 
-  int64_t resultNumber = result.getResultNumber();
-  (void)resultNumber;
-  LDBG('\n');
-  LDBG("Inplace analysis for <- #" << resultNumber << " -> #"
-                                   << operand.getOperandNumber() << " in "
-                                   << printValueInfo(result) << '\n');
-
   bool foundInterference =
       wouldCreateWriteToNonWritableBuffer(operand, result, aliasInfo, state) ||
       wouldCreateReadAfterWriteInterference(operand, result, domInfo, state,
                                             aliasInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(result);
   else
     aliasInfo.bufferizeInPlace(result, operand);
 
-  LDBG("Done inplace analysis for result #" << resultNumber << '\n');
-
   return success();
 }
 
 /// Analyze the `ops` to determine which OpResults are inplaceable. Walk ops in
 /// reverse and bufferize ops greedily. This is a good starter heuristic.
 ///
 /// Even if an op does not read or write, it may still create an alias when
 /// bufferized in-place. An example of such ops is tensor.extract_slice.

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

 //===- ModuleBufferization.cpp - Bufferization across Func. Boundaries ----===//
 //
 // 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 "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Operation.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/FormatVariadic.h"
-
-#define DEBUG_TYPE "comprehensive-module-bufferize"
-#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
-#define LDBG(X) LLVM_DEBUG(DBGS() << X)
 
 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 using namespace comprehensive_bufferize;
 
 namespace {
 /// Extra bufferization state that is required for bufferization of function
static FunctionType getOrCreateBufferizedFunctionType(
     DenseMap<FuncOp, FunctionType> &bufferizedFunctionTypes) {
   auto it = bufferizedFunctionTypes.find(funcOp);
   if (it != bufferizedFunctionTypes.end())
     return it->second;
 
   auto it2 = bufferizedFunctionTypes.try_emplace(
       funcOp, getBufferizedFunctionType(funcOp.getContext(), argumentTypes,
                                         resultTypes));
-  LDBG("FT: " << funcOp.getType() << " -> " << it2.first->second << "\n");
   return it2.first->second;
 }
 
 /// Gather equivalence info of CallOps.
 /// Note: This only adds new equivalence info if `funcOp` was already analyzed.
 // TODO: This does not handle cyclic function call graphs etc.
 static void equivalenceAnalysis(FuncOp funcOp,
                                 BufferizationAliasInfo &aliasInfo,

 /// make them inplaceable. Reliance on this logic will need to be relaxed in the
 /// future.
 ///
 /// Note: Returning a memref currently fails bufferization. If such memrefs
 /// originate from an op with an Alloc effect, they could be hoisted in the
 /// future.
 static LogicalResult bufferizeFuncOpBoundary(FuncOp funcOp,
                                              BufferizationState &state) {
-  LLVM_DEBUG(DBGS() << "Begin bufferizeFuncOpBoundary:\n" << funcOp << "\n");
   ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
 
   // If nothing to do then we are done.
   if (!llvm::any_of(funcOp.getType().getInputs(), isaTensor) &&
       !llvm::any_of(funcOp.getType().getResults(), isaTensor))
     return success();
 
   // Get the bufferized FunctionType for funcOp or construct it if not yet
static LogicalResult bufferizeFuncOpBoundary(FuncOp funcOp,
   if (funcOp.getBody().empty()) {
     if (llvm::any_of(funcOp.getType().getResults(), isaTensor))
       return funcOp->emitError() << "cannot bufferize bodiless function that "
                                  << "returns a tensor";
     FunctionType bufferizedFuncType = getOrCreateBufferizedFunctionType(
         funcOp, funcOp.getType().getInputs(), TypeRange{},
         moduleState.bufferizedFunctionTypes);
     funcOp.setType(bufferizedFuncType);
-    LLVM_DEBUG(DBGS() << "End bufferizeFuncOpBoundary no fun body: " << funcOp);
     return success();
   }
 
   // Support only single return-terminated block in the function.
   ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
   assert(returnOp && "expected func with single return op");
 
   // 1. For each FuncOp result, keep track of which inplace argument it reuses.
for (unsigned idx = 0; idx < numArgs; ++idx) {
     }
     frontBlock.eraseArgument(0);
     // TODO: add support to erase aliasInfo entries if deemed necessary.
   }
 
   // 4. Rewrite the FuncOp type to buffer form.
   funcOp.setType(bufferizedFuncType);
 
-  LLVM_DEBUG(DBGS() << "End bufferizeFuncOpBoundary:\n" << funcOp);
-
   return success();
 }
 
 /// Store all functions of the `moduleOp` in `orderedFuncOps`, sorted by
 /// callee-caller order (i.e. callees without callers first).
 /// Store the map of FuncOp to all its callers in `callerMap`.
 /// Return `failure()` if a cycle of calls is detected or if we are unable to
 /// retrieve the called FuncOp from any CallOpInterface.

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

options.allocationFns->memCpyFn = [](OpBuilder &b, Location loc, Value from,
                                        Value to) {
     b.create<linalg::CopyOp>(loc, from, to);
   };
 
   options.allowReturnMemref = allowReturnMemref;
   options.allowUnknownOps = allowUnknownOps;
   options.analysisFuzzerSeed = analysisFuzzerSeed;
   options.testAnalysisOnly = testAnalysisOnly;
+  options.printConflicts = printConflicts;
 
   // Enable InitTensorOp elimination.
   options.addPostAnalysisStep<
       linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   // TODO: Find a way to enable this step automatically when bufferizing tensor
   // dialect ops.
   options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
   options.addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();