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

Author: cxy-1993

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

@@ -1103,28 +1103,39 @@ 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) {
+  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(), 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());
   }
 }
 
 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 +1484,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 +1554,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 +1840,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 +1979,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 +2509,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());
 }
 

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";