diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h @@ -32,6 +32,9 @@ /// 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 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td @@ -297,6 +297,8 @@ "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 --- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp @@ -208,6 +208,7 @@ 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 --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp @@ -79,6 +79,8 @@ /// 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 @@ }); } +static void annotateOpsWithAliasSets(Operation *op, + const OneShotAnalysisState &state) { + AsmState asmState(op); + Builder b(op->getContext()); + op->walk([&](Operation *op) { + SmallVector aliasSets; + for (OpResult opResult : op->getOpResults()) { + if (opResult.getType().isa()) { + SmallVector 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 @@ // 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 --- 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 @@ %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