[mlir][bufferization] Add option to dump alias sets

[springerm]

Summary

This is useful for debugging.

Depends On: D143313

[springerm]

add test

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

struct OneShotBufferizationOptions : public BufferizationOptions {
   enum class AnalysisHeuristic { BottomUp, TopDown };
 
   OneShotBufferizationOptions() = default;
 
   /// Specifies whether returning newly allocated memrefs should be allowed.
   /// Otherwise, a pass failure is triggered.
   bool allowReturnAllocs = false;
 
+  /// Specifies whether the tensor IR should be annotated with alias sets.
+  bool dumpAliasSets = false;
+
   /// The heuristic controls the order in which ops are traversed during the
   /// analysis.
   AnalysisHeuristic analysisHeuristic = AnalysisHeuristic::BottomUp;
 };
 
 /// State for analysis-enabled bufferization. This class keeps track of alias
 /// to decide if tensor OpOperands should bufferize in-place.
 class OneShotAnalysisState : public AnalysisState {

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

Option<"bufferizeFunctionBoundaries", "bufferize-function-boundaries",
            "Bufferize function boundaries (experimental).">,
     Option<"copyBeforeWrite", "copy-before-write", "bool", /*default=*/"false",
            "Skip the analysis. Make a buffer copy on every write.">,
     Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
            "Specify if buffers should be deallocated. For compatibility with "
            "core bufferization passes.">,
     ListOption<"dialectFilter", "dialect-filter", "std::string",
                "Restrict bufferization to ops from these dialects.">,
+    Option<"dumpAliasSets", "dump-alias-sets", "bool", /*default=*/"false",
+           "Test only: Annotate tensor IR with alias sets">,
     ListOption<"noAnalysisFuncFilter", "no-analysis-func-filter", "std::string",
                "Skip analysis of functions with these symbol names."
                "Set copyBeforeWrite to true when bufferizing them.">,
     Option<"functionBoundaryTypeConversion",
            "function-boundary-type-conversion", "std::string",
            /*default=*/"\"infer-layout-map\"",
            "Controls layout maps when bufferizing function signatures.">,
     Option<"mustInferMemorySpace", "must-infer-memory-space", "bool",

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

if (!options) {
       // Make new bufferization options if none were provided when creating the
       // pass.
       opt.allowReturnAllocs = allowReturnAllocs;
       opt.allowUnknownOps = allowUnknownOps;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.analysisHeuristic = parseHeuristicOption(analysisHeuristic);
       opt.copyBeforeWrite = copyBeforeWrite;
       opt.createDeallocs = createDeallocs;
+      opt.dumpAliasSets = dumpAliasSets;
       opt.functionBoundaryTypeConversion =
           parseLayoutMapOption(functionBoundaryTypeConversion);
       if (mustInferMemorySpace)
         opt.defaultMemorySpace = std::nullopt;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
 

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

 // in OneShotBufferizationState. When run with `testAnalysisOnly`, the IR is
 // annotated with the results of the analysis, so that they can be checked in
 // tests.
 //===----------------------------------------------------------------------===//
 
 /// Attribute marker to specify op operands that bufferize in-place.
 constexpr StringLiteral kInPlaceOperandsAttrName = "__inplace_operands_attr__";
 
+constexpr StringLiteral kAliasSetAttrName = "__alias_set_attr__";
+
 /// Mark whether OpOperand will be bufferized inplace.
 static void setInPlaceOpOperand(OpOperand &opOperand, bool inPlace) {
   Operation *op = opOperand.getOwner();
   SmallVector<StringRef> inPlaceVector;
   if (auto attr = op->getAttr(kInPlaceOperandsAttrName)) {
     inPlaceVector = SmallVector<StringRef>(llvm::to_vector<4>(
         attr.cast<ArrayAttr>().getAsValueRange<StringAttr>()));
   } else {
annotateOpsWithBufferizationMarkers(Operation *op,
   // Add __inplace_operands_attr__.
   op->walk([&](Operation *op) {
     for (OpOperand &opOperand : op->getOpOperands())
       if (opOperand.get().getType().isa<TensorType>())
         setInPlaceOpOperand(opOperand, state.isInPlace(opOperand));
   });
 }
 
+static void annotateOpsWithAliasSets(Operation *op,
+                                     const OneShotAnalysisState &state) {
+  AsmState asmState(op);
+  Builder b(op->getContext());
+  op->walk([&](Operation *op) {
+    SmallVector<Attribute> aliasSets;
+    for (OpResult opResult : op->getOpResults()) {
+      if (opResult.getType().isa<TensorType>()) {
+        SmallVector<Attribute> aliases;
+        state.applyOnAliases(opResult, [&](Value alias) {
+          std::string buffer;
+          llvm::raw_string_ostream stream(buffer);
+          alias.printAsOperand(stream, asmState);
+          aliases.push_back(b.getStringAttr(stream.str()));
+        });
+        aliasSets.push_back(b.getArrayAttr(aliases));
+      }
+    }
+    if (!aliasSets.empty())
+      op->setAttr(kAliasSetAttrName, b.getArrayAttr(aliasSets));
+  });
+}
+
 /// Assert that every allocation can be deallocated in the same block. I.e.,
 /// every value that is returned or yielded from a block is:
 /// * guaranteed to be aliasing a bbArg of that block or a parent block, or
 /// * guaranteed to be aliasing an OpResult of a op in a parent block.
 ///
 /// In that case, buffer deallocation is simple: Every allocated buffer can be
 /// deallocated in the same block. Otherwise, the buffer deallocation pass must
 /// be run.
op->walk([&](Operation *op) {
     if (BufferizableOpInterface bufferizableOp =
             options.dynCastBufferizableOp(op))
       failedAnalysis |= failed(bufferizableOp.verifyAnalysis(state));
   });
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, state);
+  if (options.dumpAliasSets)
+    annotateOpsWithAliasSets(op, state);
 
   return success(!failedAnalysis);
 }
 
 LogicalResult
 bufferization::runOneShotBufferize(Operation *op,
                                    const OneShotBufferizationOptions &options,
                                    BufferizationStatistics *statistics) {

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir

-// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only" -allow-unregistered-dialect -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only" \
+// RUN:     -allow-unregistered-dialect -split-input-file | FileCheck %s
+
+// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only dump-alias-sets" \
+// RUN:     -allow-unregistered-dialect -split-input-file | \
+// RUN: FileCheck %s --check-prefix=CHECK-ALIAS-SETS
 
 // CHECK-LABEL: func @unknown_op_aliasing(
 func.func @unknown_op_aliasing(%f: f32, %f2: f32, %pos: index) -> f32 {
+  // CHECK-ALIAS-SETS: %[[empty:.*]] = tensor.empty
+
   %0 = tensor.empty() : tensor<10xf32>
   // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
+  // CHECK-ALIAS-SETS: %[[fill1:.*]] = linalg.fill
   %1 = linalg.fill ins(%f : f32) outs(%0 : tensor<10xf32>) -> tensor<10xf32>
 
   // Something must bufferize out-of-place because the op may return an alias
   // of %1.
   // CHECK: "dummy.dummy_op"(%{{.*}}) {__inplace_operands_attr__ = ["false"]}
   %alias = "dummy.dummy_op"(%1) : (tensor<10xf32>) -> (tensor<10xf32>)
 
   // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
+  // CHECK-ALIAS-SETS: %[[fill2:.*]] = linalg.fill {__alias_set_attr__ = [
+  // CHECK-ALIAS-SETS-SAME: ["%[[fill2]]", "%[[fill1]]", "%[[empty]]"]]
   %2 = linalg.fill ins(%f2 : f32) outs(%1 : tensor<10xf32>) -> tensor<10xf32>
   %3 = tensor.extract %alias[%pos] : tensor<10xf32>
   return %3 : f32
 }
 
 // -----
 
 // CHECK-LABEL: func @unknown_op_writing(