[mlir][linalg] Add more precise memory effects to linalg op #92079

cxy-1993 · 2024-05-14T08:10:56Z

This patch add more precise memory effect to linalg op. Including the following points:

Remove the read side effects for operands that are not used.
Set the effect for all side effects to "full".

llvmbot · 2024-05-14T08:11:27Z

@llvm/pr-subscribers-mlir-linalg
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: donald chen (cxy-1993)

Changes

This patch add more precise memory effect to linalg op. Including the following points:

Remove the read side effects for operands that are not used.
Set the effect for all side effects to "full".

Full diff: https://github.com/llvm/llvm-project/pull/92079.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+26-12)
(modified) mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp (+1-1)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index e5f83331baf81..5958e1a0f3206 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1103,20 +1103,28 @@ ParseResult GenericOp::parse(OpAsmParser &parser, OperationState &result) {
 static void getGenericEffectsImpl(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects,
-    ValueRange results, const ValueRange inputOperands,
+    LinalgOp linalgOp, ValueRange results, const ValueRange inputOperands,
     ValueRange outputOperands) {
   for (auto operand : inputOperands) {
     if (!llvm::isa<MemRefType>(operand.getType()))
       continue;
-    effects.emplace_back(MemoryEffects::Read::get(), operand,
+    effects.emplace_back(MemoryEffects::Read::get(), 0, true, operand,
                          SideEffects::DefaultResource::get());
   }
-  for (auto operand : outputOperands) {
+  unsigned inputOperandSize = inputOperands.size();
+  unsigned usedOutputSize =
+      linalgOp.getOpOperandsMatchingBBargs().size() - inputOperandSize;
+
+  for (auto [index, operand] : llvm::enumerate(outputOperands)) {
     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 (index < usedOutputSize &&
+        linalgOp.payloadUsesValueFromOperand(
+            &linalgOp->getOpOperand(index + inputOperandSize))) {
+      effects.emplace_back(MemoryEffects::Read::get(), 0, true, operand,
+                           SideEffects::DefaultResource::get());
+    }
+    effects.emplace_back(MemoryEffects::Write::get(), 0, true, operand,
                          SideEffects::DefaultResource::get());
   }
 }
@@ -1124,7 +1132,8 @@ static void getGenericEffectsImpl(
 void GenericOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
-  getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
+  getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()),
+                        getOperation()->getResults(), getDpsInputs(),
                         getDpsInits());
 }
 
@@ -1473,7 +1482,8 @@ ArrayAttr MapOp::getIndexingMaps() {
 void MapOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
-  getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
+  getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()),
+                        getOperation()->getResults(), getDpsInputs(),
                         getDpsInits());
 }
 
@@ -1542,7 +1552,8 @@ ArrayAttr ReduceOp::getIndexingMaps() {
 void ReduceOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
-  getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
+  getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()),
+                        getOperation()->getResults(), getDpsInputs(),
                         getDpsInits());
 }
 
@@ -1827,7 +1838,8 @@ ArrayAttr TransposeOp::getIndexingMaps() {
 void TransposeOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
-  getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
+  getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()),
+                        getOperation()->getResults(), getDpsInputs(),
                         getDpsInits());
 }
 
@@ -1965,7 +1977,8 @@ ArrayAttr BroadcastOp::getIndexingMaps() {
 void BroadcastOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
-  getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
+  getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()),
+                        getOperation()->getResults(), getDpsInputs(),
                         getDpsInits());
 }
 
@@ -2494,7 +2507,8 @@ SoftmaxOp::reifyResultShapes(OpBuilder &b,
 void SoftmaxOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
-  getGenericEffectsImpl(effects, getOperation()->getResults(), getDpsInputs(),
+  getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()),
+                        getOperation()->getResults(), getDpsInputs(),
                         getDpsInits());
 }
 
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..f3071b81e21cb 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,7 +659,7 @@ LogicalResult {0}::fold(FoldAdaptor,
 void {0}::getEffects(SmallVectorImpl<
     SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
       if (hasPureTensorSemantics()) return;
-      getGenericEffectsImpl(effects,
+      getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()),
         getOperation()->getResults(), getDpsInputs(), getDpsInits());
 }
 )FMT";

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

cxy-1993 · 2024-05-16T01:59:40Z

Counld you please help me merge this patch into master? @matthias-springer
After your last review, I made changes to the implementation of payloadUsesValueFromOperand. The previous implementation could cause a core dump in some ops.

ftynse · 2024-05-17T11:50:15Z

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

+      if (getOperation()->getRegion(0).empty()) {
+        return true;
+      }


I don't quite get this. When the region is empty (is this even allowed?), the operand is considered used? But when it has a terminator not reading from the corresponding block argument, the operand is not considered used? Why? If we assume it is read anyway, it should be marked as such in both cases.

I don't quite get this. When the region is empty (is this even allowed?), the operand is considered used? But when it has a terminator not reading from the corresponding block argument, the operand is not considered used? Why? If we assume it is read anyway, it should be marked as such in both cases.

Thank you for your suggestion! For this issue, here is my understanding: When region is empty, it usually occurs in scenarios where region does not need to interpret the semantics of op, such as linalg.map:

%mapped = linalg.map { arith.addf } ins(%arg0, %arg1 : tensor<10x100xf32>, tensor<10x100xf32>) outs(%map_init : tensor<10x100xf32>)

In the example above, the linalg.map has an empty region. This is because, with arith.add present, we do not need to provide additional interpretation for this op, as its implementation simply involves adding the two input operands. In such scenarios where no interpretation is required, we consider all operands to be used in accordance with the semantics of the op.

In linalg operations implemented with regions, it is evident that when the block arguments corresponding to operands are not used, it indicates that they are not being read.

linalg.map { arith.addf } is just a pretty-printed version. It has a region with an arith.addf and linalg.yield terminator.

linalg.map { arith.addf } is just a pretty-printed version. It has a region with an arith.addf and linalg.yield terminator.

You are right. I reviewed why I thought this way: before this, I didn't add handling for when region was empty, leading to some core dump issues during part of the bufferization tests. This happened because the 'cloneWithoutRegions' function was called during the bufferization process, and in this function, the side effects of 'op' were queried, unfortunately, at this time region was still empty. So, I'd like to solicit opinions from both of you: in this situation, can we conservatively assume that all operands can be read and written to? @matthias-springer @ftynse

Hi @matthias-springer @ftynse,I wanted to send a friendly reminder to re-review the proposal above at your earliest convenience and hope to hear your thoughts on this issue.

Why does this lead to a crash? The bufferization impl for LinalgOps calls cloneWithoutRegions, but inlines the original region immediately afterwards. There should be no call to payloadUsesValueFromOperand in-between. (Even if there is one, this function would just return false and should not crash.)

Basically, the op is invalid without a region body. Generally speaking, helper functions on an op are allowed to function incorrectly if the op is invalid.

Why does this lead to a crash?

When the cloneWithoutRegions function is called, the bufferize rewriter will invoke the notifyOperationInserted function. This function will query the side effects of the operation to count the number of memory allocations. Because this patch will call payloadUsesValueFromOperand when querying side effects, and the payloadUsesValueFromOperand of the map operation is implemented as follows:

bool payloadUsesValueFromOperand(OpOperand * opOperand) { if (isDpsInit(opOperand)) return false; return !getMatchingBlockArgument(opOperand).use_empty(); }

This will fetch the block inside the op's region, causing a core dump.

Basically, the op is invalid without a region body. Generally speaking, helper functions on an op are allowed to function incorrectly if the op is invalid.

So should we prohibit querying side effects before the block is inserted?

Oh, I see. The op is already inserted even though it was not fully constructed yet (region missing). Ideally, we would change the BufferizableOpInterface impl., so that we insert the op only when it is fully constructed.

I think it's important to keep the LinalgOp interface clean and we should not have to account for edge cases in the bufferization.

Could you change the implementation in Linalg/Transforms/BufferizableOpInterfaceImpl.cpp such that it clones the op without regions without a builder. Basically using Operation::create(OpState) (instead of cloneOpWithoutRegions). Then inline the region. Then rewriter.insert. Then the callback in the bufferization will be triggered on a fully constructed and valid op.

I agree with your perspective, we should keep the implementation within the interface clean. I have already modified the code according to your comment, thanks for your feedback.

Hi @ftynse, do you have any further comments on the modifications proposed in this patch?

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

cxy-1993 · 2024-05-28T06:04:14Z

Our friend ftynse hasn't had time to work on this patch recently, and the issues he raised seem to have been resolved. Could you please help me merge this PR? @matthias-springer

llvmbot added mlir:core MLIR Core Infrastructure mlir:linalg mlir labels May 14, 2024

cxy-1993 force-pushed the linalg-effect branch 2 times, most recently from 73f84df to a3bd52c Compare May 14, 2024 08:14

matthias-springer reviewed May 14, 2024

View reviewed changes

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp Outdated Show resolved Hide resolved

cxy-1993 force-pushed the linalg-effect branch from a3bd52c to 9d40a4f Compare May 14, 2024 10:38

cxy-1993 requested a review from matthias-springer May 14, 2024 10:38

matthias-springer reviewed May 14, 2024

View reviewed changes

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp Outdated Show resolved Hide resolved

cxy-1993 force-pushed the linalg-effect branch 2 times, most recently from 71d6a49 to 8d3b0c8 Compare May 14, 2024 12:31

cxy-1993 requested a review from matthias-springer May 14, 2024 12:33

matthias-springer reviewed May 14, 2024

View reviewed changes

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp Outdated Show resolved Hide resolved

cxy-1993 requested a review from matthias-springer May 14, 2024 12:50

matthias-springer approved these changes May 14, 2024

View reviewed changes

cxy-1993 force-pushed the linalg-effect branch 4 times, most recently from 779df1f to 9ccde8b Compare May 15, 2024 06:37

cxy-1993 requested a review from matthias-springer May 16, 2024 03:35

ftynse requested changes May 17, 2024

View reviewed changes

cxy-1993 force-pushed the linalg-effect branch from 9ccde8b to 67321b7 Compare May 17, 2024 13:17

cxy-1993 requested a review from ftynse May 17, 2024 13:28

[mlir][linalg] Add more precise memory effects to linalg op

43ed2f0

cxy-1993 force-pushed the linalg-effect branch from 67321b7 to 43ed2f0 Compare May 23, 2024 14:52

matthias-springer approved these changes May 24, 2024

View reviewed changes

matthias-springer merged commit a0fdb38 into llvm:main May 28, 2024
3 checks passed

LongshengDu mentioned this pull request Jul 1, 2024

Update llvm version for pack propagation and linalg memory effects helpers intel/graph-compiler#136

Merged

[mlir][linalg] Add more precise memory effects to linalg op #92079

[mlir][linalg] Add more precise memory effects to linalg op #92079

Uh oh!

Conversation

cxy-1993 commented May 14, 2024

Uh oh!

llvmbot commented May 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cxy-1993 commented May 16, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

matthias-springer May 23, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

cxy-1993 commented May 28, 2024

Uh oh!

Uh oh!

Uh oh!

llvmbot commented May 14, 2024 •

edited

Loading

cxy-1993 commented May 16, 2024 •

edited

Loading

matthias-springer May 23, 2024 •

edited

Loading