[Mlir-commits] [mlir] [mlir][linalg] Add more precise memory effects to linalg op (PR #92079)
donald chen
llvmlistbot at llvm.org
Thu May 23 07:52:06 PDT 2024
https://github.com/cxy-1993 updated https://github.com/llvm/llvm-project/pull/92079
>From 43ed2f0953fae9dd6bb246f6300e31ac3e63a34b Mon Sep 17 00:00:00 2001
From: cxy <chenxunyu1993 at gmail.com>
Date: Tue, 14 May 2024 08:03:42 +0000
Subject: [PATCH] [mlir][linalg] Add more precise memory effects to linalg op
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 62 ++++++++++++-------
.../BufferizableOpInterfaceImpl.cpp | 14 +++--
.../mlir-linalg-ods-yaml-gen.cpp | 3 +-
3 files changed, 52 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index e5f83331baf81..3cdd606297091 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1103,20 +1103,30 @@ ParseResult GenericOp::parse(OpAsmParser &parser, OperationState &result) {
static void getGenericEffectsImpl(
SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
&effects,
- ValueRange results, const ValueRange inputOperands,
- ValueRange outputOperands) {
- for (auto operand : inputOperands) {
+ LinalgOp linalgOp) {
+ SmallVector<Value> inputOperands = linalgOp.getDpsInputs();
+ for (auto [index, operand] : llvm::enumerate(inputOperands)) {
if (!llvm::isa<MemRefType>(operand.getType()))
continue;
- effects.emplace_back(MemoryEffects::Read::get(), operand,
- SideEffects::DefaultResource::get());
+ if (linalgOp.payloadUsesValueFromOperand(&linalgOp->getOpOperand(index))) {
+ effects.emplace_back(MemoryEffects::Read::get(), operand, /*stage=*/0,
+ /*effectOnFullRegion=*/true,
+ SideEffects::DefaultResource::get());
+ }
}
- for (auto operand : outputOperands) {
+ unsigned inputOperandSize = inputOperands.size();
+
+ for (auto [index, operand] : llvm::enumerate(linalgOp.getDpsInits())) {
if (!llvm::isa<MemRefType>(operand.getType()))
continue;
- effects.emplace_back(MemoryEffects::Read::get(), operand,
- SideEffects::DefaultResource::get());
- effects.emplace_back(MemoryEffects::Write::get(), operand,
+ if (linalgOp.payloadUsesValueFromOperand(
+ &linalgOp->getOpOperand(index + inputOperandSize))) {
+ effects.emplace_back(MemoryEffects::Read::get(), operand, /*stage=*/0,
+ /*effectOnFullRegion=*/true,
+ SideEffects::DefaultResource::get());
+ }
+ effects.emplace_back(MemoryEffects::Write::get(), operand, /*stage=*/0,
+ /*effectOnFullRegion=*/true,
SideEffects::DefaultResource::get());
}
}
@@ -1124,8 +1134,7 @@ static void getGenericEffectsImpl(
void GenericOp::getEffects(
SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
&effects) {
- getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
- getDpsInits());
+ getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
}
LogicalResult GenericOp::verify() { return success(); }
@@ -1473,8 +1482,7 @@ ArrayAttr MapOp::getIndexingMaps() {
void MapOp::getEffects(
SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
&effects) {
- getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
- getDpsInits());
+ getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
}
//===----------------------------------------------------------------------===//
@@ -1542,8 +1550,7 @@ ArrayAttr ReduceOp::getIndexingMaps() {
void ReduceOp::getEffects(
SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
&effects) {
- getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
- getDpsInits());
+ getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
}
static ParseResult parseDenseI64ArrayAttr(OpAsmParser &parser,
@@ -1827,8 +1834,7 @@ ArrayAttr TransposeOp::getIndexingMaps() {
void TransposeOp::getEffects(
SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
&effects) {
- getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
- getDpsInits());
+ getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
}
LogicalResult TransposeOp::fold(FoldAdaptor adaptor,
@@ -1965,8 +1971,7 @@ ArrayAttr BroadcastOp::getIndexingMaps() {
void BroadcastOp::getEffects(
SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
&effects) {
- getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
- getDpsInits());
+ getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
}
void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
@@ -2494,8 +2499,23 @@ SoftmaxOp::reifyResultShapes(OpBuilder &b,
void SoftmaxOp::getEffects(
SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
&effects) {
- getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
- getDpsInits());
+ for (Value operand : getDpsInputs()) {
+ if (!llvm::isa<MemRefType>(operand.getType()))
+ continue;
+ effects.emplace_back(MemoryEffects::Read::get(), operand, /*stage=*/0,
+ /*effectOnFullRegion=*/true,
+ SideEffects::DefaultResource::get());
+ }
+ for (Value operand : getDpsInits()) {
+ if (!llvm::isa<MemRefType>(operand.getType()))
+ continue;
+ effects.emplace_back(MemoryEffects::Read::get(), operand, /*stage=*/0,
+ /*effectOnFullRegion=*/true,
+ SideEffects::DefaultResource::get());
+ effects.emplace_back(MemoryEffects::Write::get(), operand, /*stage=*/0,
+ /*effectOnFullRegion=*/true,
+ SideEffects::DefaultResource::get());
+ }
}
// Helper functions for softmax decomposition.
diff --git a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
index 899b8c87d0df7..81a5398dabcb7 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -76,10 +76,16 @@ bufferizeDestinationStyleOpInterface(RewriterBase &rewriter,
// new op. Since the new op does not have any tensor results, it does not
// return anything.
assert(op->getNumRegions() == 1 && "expected that op has 1 region");
- auto newOp = cast<DestinationStyleOpInterface>(cloneWithoutRegions(
- rewriter, op, /*newResultTypes=*/TypeRange{}, newOperands));
- rewriter.inlineRegionBefore(op->getRegion(0), newOp->getRegion(0),
- newOp->getRegion(0).begin());
+ OperationState state(op->getLoc(), op->getName(), newOperands, TypeRange{},
+ op->getAttrs());
+ state.addRegion();
+ Operation *newOp = Operation::create(state);
+ newOp->getRegion(0).getBlocks().splice(newOp->getRegion(0).begin(),
+ op->getRegion(0).getBlocks());
+
+ // We don't want the rewriter tracks an incomplete operation, so insert new
+ // operation after op was fully constructed.
+ rewriter.insert(newOp);
// Replace the results of the old op with the new output buffers.
replaceOpWithBufferizedValues(rewriter, op, newOutputBuffers);
diff --git a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
index fe6ad15041126..882000ee0969a 100644
--- a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
+++ b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
@@ -659,8 +659,7 @@ LogicalResult {0}::fold(FoldAdaptor,
void {0}::getEffects(SmallVectorImpl<
SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
if (hasPureTensorSemantics()) return;
- getGenericEffectsImpl(effects,
- getOperation()->getResults(), getDpsInputs(), getDpsInits());
+ getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
}
)FMT";
More information about the Mlir-commits
mailing list