[mlir][linalg] Remove unused payload related OutOpOperand

[raikonenfnu]

Summary

Some higher level operations such as torch.max generates linalg generic
that returns both the index and the value of the max operation. However
sometimes not all information is being used. This however blocks
vectorization for certain cases which causes performance degradation.
This patch aims to fix this issue.

[raikonenfnu]

Added missing check

[mravishankar]

Thanks Stanley! Have a few refactoring requests to make the code more understandable.

[mravishankar]

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

1034

No related to this patch, but just to make it more readable since the method is getting too big now. Can you change this to

if (!generic.hasTensorSemantics()) {
  ....
  return dedupedOutputs;
}
/// body of the else goes here....

1061–1063

To make the logic easier to follow do you mind moving lines 1006 - 1017 and 1030 - 1074 into a helper function that just uses the outputOpOperand.

So something like

bool isOutputOpOperandDead(linalg::GenericOp genericOp, OpOperand &outOperand, OpResult result) {
  if (!result.use_empty())
    return false;
  if(!genericOp.payloadUsesValueFromOperand(outputOpOperand))
    return true;
  // Logic for cycle added in this change.
}

... deduplicateOutputOperands ... {
 ...
 for (const auto &outputOperand : ...) {
   ....
   if (isOutputOpOperandDead(...)) {
     droppedOpOperands.push_back(...)
     if (genericOp.canOpOperandsBeDropped(droppedOpOperands))
        continue;
     droppedOpOperands.pop_back();
   }

   // Lines 1024 - 1029....
}

If this wasnt easy to follow, feel free to punt to me.

[raikonenfnu]

Addressed refactoring comments.

[raikonenfnu]

Test to be added

In D135388#3841077, @mravishankar wrote:

Thanks Stanley! Have a few refactoring requests to make the code more understandable.

Thanks for the refactoring comments and suggestions Mahesh! :)

[mravishankar]

THanks for refactoring. More comments. Marking as requesting change till tests are added.

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

867

Nit: I think you can use genericOp.getBody().

869

This can be an assert I think

873

Please use getRegionOutputArgs (https://github.com/llvm/llvm-project/blob/3b20765cf725e57d0f99c3d292fe4891b0ee801d/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td#L364) and get outArg from there.

Even better, just take the BlockArgument as the argument instead of OpIndex.

[raikonenfnu]

Rebased on later llvm

[raikonenfnu]

Addressed comments/cleaned up with block args

[raikonenfnu]

Exclude dynamic case which was unsupported to begin with and is unlikely to be outs anyways.

[mravishankar]

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

1109

I don't think we need to check this. The dynamic case is not unsupported. It is expected to be supported. Could you provide more context here?

[raikonenfnu]

@mravishankar It was breaking one of IREE's e2e test https://github.com/iree-org/iree/blob/main/tests/e2e/regression/dynamic_torch_index_select_negative.mlir. It tried creating a placeholder/dummy tensor for a dynamic tensor and it broke.

[ThomasRaoux]

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

876–886

the operand index of the yield has to match block argument index otherwise it is wrong.

it could be something like:
^bb(%arg0, %arg1)

yield %arg1, %arg0

[raikonenfnu]

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

876–886

shouldn't this not matter for this part of the code since, this part is only checking for uses?

[ThomasRaoux]

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

876–886

it's mean to check if the cycle is dead right? Doesn't this affect it?

[raikonenfnu]

Addressed indexing mismatch issue and separated unused cycle as it's own pattern.

[raikonenfnu]

Minor comment change.

[raikonenfnu]

Added some test.

[raikonenfnu]

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

876–886

Understood, Mahesh also helped clarify this, thanks for this suggestion :)

[raikonenfnu]

Updated to handle genericOp with in place updates.
Look at tile-parallel-reduce.mlir

[raikonenfnu]

minor comment change

[raikonenfnu]

Add more indexing check

[raikonenfnu]

Added test where some cycles/bout are used by others outArgs.

[raikonenfnu]

minor nit on test comments.

[raikonenfnu]

nit clean up

[mravishankar]

Mostly looks good. Have a few nits. Should be ready to land after that!

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

852

You can drop BlockArgument & for BlockArgument .

Could we make the ABI a little simpler (we want to avoid the implicit assumption that outputs come after inputs in linalg ops). So make this

bool isResultValueDead(linalg::GenericOp genericOp, OpResult result).

Then

outputOpOperand = genericOp.getOutputOperands(result.getResultNumber());
outputArg = genericOp.getRegionOutputArgs(result.getResultNumber());

Also below to check the usage in the yield you can just do

if (yieldOp.getOperand(result.getResultNumber()) != outputArg) return false;

871

You can drop this check. You are checking below argUserOp is a linalg::YieldOp. That has no uses by construction.

876

Replace these lines with

auto yieldOp = dyn_cast<linalg::YieldOp>(argUserOp);
if (!yieldOp) return false;

Then see above for changing the condition check.

1111

Do we still need to do this?

1240

I think you can replace this with

if (!genericOp.hasTensorSemantics())
  return failure()

Also note, LLVM style is to not have braces around bodies that are a single statement (unlike IREE for example)

[ThomasRaoux]

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

851

this should be a static function

1277

What is it meant to do? I don't think it does anything?

[raikonenfnu]

Clean up and addressed comments.

[raikonenfnu]

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

1111

we still need it since before dropping the yield it does rewriter.mergeBlocks which will would cause segfault otherwise

1277

I think it's to in

1277

Don't we need the return success to indicate matchAndRewrite is successful?

[ThomasRaoux]

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

1277

looks like my comment shifted line after you latest update. My comment was about this line:
rewriter.updateRootInPlace(genericOp, [] {});

[raikonenfnu]

Addressed Thomas comments and suggetions

[mravishankar]

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

1276

Instead of returning from within the loop, can you just check if a modification was made or not in the loop and return success or failure after the loop

[raikonenfnu]

Moved return out of the loop, brought back updateRootInPlace, and moved yieldOp to before mergeblocks s.t we don't need to replace blockarguments

[raikonenfnu]

minor style change

[mravishankar]

Thanks!

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

1281

Nit: return success(hasRemovedCycles);

[stella.stamenova]

Looks like this broke the windows mlir buildbot: https://lab.llvm.org/buildbot/#/builders/13/builds/26882

[raikonenfnu]

In D135388#3847663, @stella.stamenova wrote:

Looks like this broke the windows mlir buildbot: https://lab.llvm.org/buildbot/#/builders/13/builds/26882

Hey @stella.stamenova, thanks for the notifications, I am pushing up a fix here https://reviews.llvm.org/D135612 :)

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

851

done, thanks!

852

done!

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

void GenericOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   SmallVector<Value> inputBuffers = getInputBufferOperands();
   SmallVector<Value> outputBuffers = getOutputBufferOperands();
   getGenericEffectsImpl(effects, getOperation()->getResults(), inputBuffers,
                         outputBuffers);
 }
 
+bool isOutputOpOperandDead(linalg::GenericOp genericOp,
+                           OpOperand *outputOpOperand, uint64_t OpIndex,
+                           Value result) {
+  if (!result.use_empty())
+    return false;
+  // If out operand not used in payload, we can drop it.
+  if (!genericOp.payloadUsesValueFromOperand(outputOpOperand))
+    return true;
+
+  // The out operand that is part of a payload can be dropped if
+  // these conditions are met:
+  // - Result from out operand is dead.
+  // - Block arg from out operand has a single use in the %cycle
+  // instruction.
+  // - Cycle has a single use in the yield or it itself is a yield.
+
+  Block *genericBlock = &genericOp.getRegion().front();
+  uint64_t argIndex = genericOp.getNumInputs() + OpIndex;
+  if (argIndex >= genericBlock->getNumArguments())
+    return false;
+
+  // Check block arg and cycle from out operand has a single use.
+  Value outArg = genericBlock->getArgument(argIndex);
+  if (!outArg.hasOneUse())
+    return false;
+  Operation *cycleOp = *outArg.user_begin();
+
+  // Check cycle has at most one use.
+  if (!cycleOp->hasOneUse() && !cycleOp->use_empty())
+    return false;
+
+  // Check that if it has one use it is a yield.
+  if (cycleOp->hasOneUse()) {
+    if (!isa<linalg::YieldOp>(*cycleOp->user_begin()))
+      return false;
+  }
+
+  // Check that if it has no use, cycle itself is a yield.
+  if (cycleOp->use_empty()) {
+    if (!isa<linalg::YieldOp>(cycleOp))
+      return false;
+  }
+
+  return true;
+}
+
 LogicalResult GenericOp::verify() { return success(); }
 
 namespace {
 
 struct DeduplicateAndRemoveDeadOperandsAndResults
     : public OpRewritePattern<GenericOp> {
   using OpRewritePattern<GenericOp>::OpRewritePattern;
 
deduplicateOutputOperands(GenericOp genericOp,
     // preserved.
     if (!genericOp.hasTensorSemantics()) {
       for (const auto &outputOpOperand :
            llvm::enumerate(genericOp.getOutputOperands())) {
         origToNewPos[outputOpOperand.index()] = newOutputOperands.size();
         newOutputOperands.push_back(outputOpOperand.value()->get());
         newIndexingMaps.push_back(
             genericOp.getMatchingIndexingMap(outputOpOperand.value()));
       }
-    } else {
-      // Output argument can be dropped if the result has
-      // - no users, and
-      // - it is not used in the payload, and
-      // - the corresponding indexing maps are not needed for loop bound
-      //   computation.
-      auto yieldOp = cast<YieldOp>(genericOp.getBody()->getTerminator());
-      for (const auto &outputOpOperand :
-           llvm::enumerate(genericOp.getOutputOperands())) {
-        Value result = genericOp.getResult(outputOpOperand.index());
-        AffineMap indexingMap =
-            genericOp.getMatchingIndexingMap(outputOpOperand.value());
-        auto key =
-            std::make_tuple(outputOpOperand.value()->get(), indexingMap,
-                            yieldOp->getOperand(outputOpOperand.index()));
+      return origToNewPos;
+    }
 
-        // Do not drop an out if its value is used in the payload.
-        if (!genericOp.payloadUsesValueFromOperand(outputOpOperand.value())) {
-          if (result.use_empty()) {
-            // Check if the opoperand can be dropped without affecting loop
-            // bound computation. Add the operand to the list of dropped op
-            // operand for checking. If it cannot be dropped, need to pop the
-            // value back.
-            droppedOpOperands.push_back(outputOpOperand.value());
-            if (genericOp.canOpOperandsBeDropped(droppedOpOperands)) {
-              continue;
-            }
-            droppedOpOperands.pop_back();
-          }
+      if (isOutputOpOperandDead(genericOp, outputOpOperand.value(),
+                              outputOpOperand.index(), result)) {
 
-          // The out operand can also be dropped if it is computed redundantly
-          // by another result, the conditions for that are
-          // - The same operand is used as the out operand
-          // - The same indexing map is used
-          // - The same yield value is used.
-          auto it = dedupedOutpts.find(key);
-          if (it != dedupedOutpts.end()) {
-            origToNewPos[outputOpOperand.index()] = it->second;
-            droppedOpOperands.push_back(outputOpOperand.value());
-            continue;
-          }
-        }
+      if (!genericOp.payloadUsesValueFromOperand(outputOpOperand.value())) {
 
-        origToNewPos[outputOpOperand.index()] = newOutputOperands.size();
-        dedupedOutpts[key] = newOutputOperands.size();
-        newOutputOperands.push_back(outputOpOperand.value()->get());
-        newIndexingMaps.push_back(
-            genericOp.getMatchingIndexingMap(outputOpOperand.value()));
-      }
-    }
-
     return origToNewPos;
   }
 
   // Populate the body of the canonicalized operation.
   void populateOpPayload(
       GenericOp genericOp, GenericOp newOp,
       const llvm::SmallDenseMap<unsigned, unsigned> &origInsToNewInsPos,
       const llvm::SmallDenseMap<unsigned, unsigned> &origOutsToNewOutsPos,
       PatternRewriter &rewriter) const {
     // Merge the body of the original op with the new op.
     Block *newOpBlock = &newOp.getRegion().front();
     assert(newOpBlock->empty() && "expected new op to have an empty payload");
     Block *origOpBlock = &genericOp.getRegion().front();
     SmallVector<Value> replacements(origOpBlock->getNumArguments(), nullptr);
 
     // Replace all arguments in the original op, with arguments from the
     // canonicalized op.
     auto updateReplacements =
         [&](OpOperandVector &origOperands, OpOperandVector &newOperands,
             const llvm::SmallDenseMap<unsigned, unsigned> &map) {
           for (const auto &origOperand : llvm::enumerate(origOperands)) {
             auto it = map.find(origOperand.index());
-            if (it == map.end())
+            if (it == map.end()) {
+              // Creates a placeholder constOp for deleted OpOperands.
+              Type elementType = origOperand.value()->get().getType();
+              Value placeHolder = rewriter.create<arith::ConstantOp>(
+                  newOpBlock->getParent()->getLoc(), elementType,
+                  rewriter.getZeroAttr(elementType));
+              replacements[origOperand.value()->getOperandNumber()] =
+                  placeHolder;
               continue;
+            }
             OpOperand *newOperand = newOperands[it->second];
             replacements[origOperand.value()->getOperandNumber()] =
                 newOpBlock->getArgument(newOperand->getOperandNumber());
           }
         };
 
     OpOperandVector origInputOperands = genericOp.getInputOperands();
     OpOperandVector newInputOperands = newOp.getInputOperands();

 }
 
 LogicalResult GenericOp::fold(ArrayRef<Attribute>,
                               SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // ReduceOp
 //===----------------------------------------------------------------------===//
 
 ArrayAttr ReduceOp::getIteratorTypes() {
   int64_t inputRank = getInputs()[0].getType().cast<ShapedType>().getRank();
   SmallVector<StringRef> iteratorTypes(inputRank,
                                        getParallelIteratorTypeName());
   for (int64_t reductionDim : getDimensions())
     iteratorTypes[reductionDim] = getReductionIteratorTypeName();
void ReduceOp::getEffects(
   SmallVector<Value> inputBuffers = getInputBufferOperands();
   SmallVector<Value> outputBuffers = getOutputBufferOperands();
   getGenericEffectsImpl(effects, getOperation()->getResults(), inputBuffers,
                         outputBuffers);
 }
 
 static ParseResult parseDstStyleOp(
     OpAsmParser &parser, OperationState &result,
     function_ref<ParseResult(OpAsmParser &, NamedAttrList &)> parseAttrsFn =
         nullptr) {
   // Parse `ins` and `outs`.
   SmallVector<Type, 4> inputTypes, outputTypes;
   if (parseCommonStructuredOpParts(parser, result, inputTypes, outputTypes,
                                    /*addOperandSegmentSizes=*/false))
     return failure();
 
   // Add result types.
   for (Type outputType : outputTypes) {
     if (!outputType.isa<RankedTensorType>())
       return failure();
     result.addTypes(outputType);
   }