[Mlir-commits] [mlir] bb9d1b5 - [mlir][bufferization] Add option to dump alias sets
Matthias Springer
llvmlistbot at llvm.org
Mon May 15 06:38:41 PDT 2023
Author: Matthias Springer
Date: 2023-05-15T15:38:20+02:00
New Revision: bb9d1b551a4407660293c2bf3f2343ba70ed8e68
URL: https://github.com/llvm/llvm-project/commit/bb9d1b551a4407660293c2bf3f2343ba70ed8e68
DIFF: https://github.com/llvm/llvm-project/commit/bb9d1b551a4407660293c2bf3f2343ba70ed8e68.diff
LOG: [mlir][bufferization] Add option to dump alias sets
This is useful for debugging.
Differential Revision: https://reviews.llvm.org/D143314
Added:
Modified:
mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
index 8a7d8f0abb5b8..4fd3da1548c41 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 20b35f89ae269..d801aec2a2dc2 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 24aaff0e48826..7766a8a32be7f 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 a9f05b21282de..5e3519669ec84 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 4eaa7dc2bcbf5..df174d01ca07a 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
More information about the Mlir-commits
mailing list