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

This is useful for debugging.

Differential Revision: https://reviews.llvm.org/D143314

5 files changed, 44 insertions, 1 deletions

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
index 8a7d8f0abb5b..4fd3da1548c4 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
@@ -32,6 +32,9 @@ struct OneShotBufferizationOptions : public BufferizationOptions {
   /// 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;
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 20b35f89ae26..d801aec2a2dc 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -297,6 +297,8 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
            "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.">,
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 24aaff0e4882..7766a8a32be7 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -208,6 +208,7 @@ struct OneShotBufferizePass
       opt.analysisHeuristic = parseHeuristicOption(analysisHeuristic);
       opt.copyBeforeWrite = copyBeforeWrite;
       opt.createDeallocs = createDeallocs;
+      opt.dumpAliasSets = dumpAliasSets;
       opt.setFunctionBoundaryTypeConversion(
           parseLayoutMapOption(functionBoundaryTypeConversion));
       if (mustInferMemorySpace)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index a9f05b21282d..5e3519669ec8 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -79,6 +79,8 @@ static bool isaTensor(Type t) { return isa<TensorType>(t); }
 /// 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();
@@ -986,6 +988,29 @@ annotateOpsWithBufferizationMarkers(Operation *op,
   });
 }
 
+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
@@ -1100,6 +1125,8 @@ LogicalResult bufferization::analyzeOp(Operation *op,
   // 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);
 }
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
index 4eaa7dc2bcbf..df174d01ca07 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
@@ -1,9 +1,17 @@
-// 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
@@ -12,6 +20,8 @@ func.func @unknown_op_aliasing(%f: f32, %f2: f32, %pos: index) -> f32 {
   %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