[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

988

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....

1030

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

1108

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.

1110

Do we still need to do this?

1239

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

1276

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

1110

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

1276

I think it's to in

1276

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

[ThomasRaoux]

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

1276

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

1275

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

1280

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);
 }
 
 LogicalResult GenericOp::verify() { return success(); }
 
 namespace {
 
 struct DeduplicateAndRemoveDeadOperandsAndResults
     : public OpRewritePattern<GenericOp> {
   using OpRewritePattern<GenericOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     // Create a map from argument position in the original op to the argument
     // position in the new op. If the argument is dropped it wont have an entry.
     SmallVector<OpOperand *> droppedOpOperands;
 
     // Information needed to build the new op.
     SmallVector<Value> newInputOperands, newOutputOperands;
     SmallVector<AffineMap> newIndexingMaps;
 
     // Gather information about duplicate input operands.
     llvm::SmallDenseMap<unsigned, unsigned> origInsToNewInsPos =
         deduplicateInputOperands(genericOp, droppedOpOperands, newInputOperands,
                                  newIndexingMaps);
 
     // Gather information about the dropped outputs.
     llvm::SmallDenseMap<unsigned, unsigned> origOutsToNewOutsPos =
         deduplicateOutputOperands(genericOp, droppedOpOperands,
                                   newOutputOperands, newIndexingMaps);
 
     // Check if there is any change to operands.
     if (newInputOperands.size() + newOutputOperands.size() ==
         static_cast<size_t>(genericOp.getNumInputsAndOutputs()))
       return failure();
 
     // Create the new op with the body being empty.
     Location loc = genericOp.getLoc();
     SmallVector<Type> newResultTypes;
     if (genericOp.hasTensorSemantics()) {
       newResultTypes = llvm::to_vector(llvm::map_range(
           newOutputOperands, [](Value v) { return v.getType(); }));
     }
     auto newOp = rewriter.create<GenericOp>(
         loc, newResultTypes, newInputOperands, newOutputOperands,
         rewriter.getAffineMapArrayAttr(newIndexingMaps),
         genericOp.getIteratorTypes(), genericOp.getDocAttr(),
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.getTiedIndexingMap(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 :
if (!genericOp.hasTensorSemantics()) {
           // - 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;
           }
+        } else {
+          // 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.
+
+          // Checking result from out operand is dead.
+          if (result.use_empty()) {
+            // Obtaining valid block argument.
+            Block *genericBlock = &genericOp.getRegion().front();
+            uint64_t argIndex =
+                genericOp.getNumInputs() + outputOpOperand.index();
+            if (argIndex >= genericBlock->getNumArguments())
+              continue;
+
+            // Check block arg and cycle from out operand has a single use.
+            Value outArg = genericBlock->getArgument(argIndex);
+            if (!outArg.hasOneUse())
+              continue;
+            Operation *cycleOp = *outArg.user_begin();
+
+            // Check cycle has at most one use.
+            if (!cycleOp->hasOneUse() && !cycleOp->use_empty())
+              continue;
+
+            // Check that if it has one use it is a yield.
+            if (cycleOp->hasOneUse()) {
+              if (!isa<linalg::YieldOp>(*outArg.user_begin()))
+                continue;
+            }
+
+            // Check that if it has no use, cycle itself is a yield.
+            if (cycleOp->use_empty()) {
+              if (!isa<linalg::YieldOp>(cycleOp))
+                continue;
             }
 
+            droppedOpOperands.push_back(outputOpOperand.value());
+            if (genericOp.canOpOperandsBeDropped(droppedOpOperands)) {
+              continue;
+            }
+            droppedOpOperands.pop_back();
+          }
+        }
         origToNewPos[outputOpOperand.index()] = newOutputOperands.size();
         dedupedOutpts[key] = newOutputOperands.size();
         newOutputOperands.push_back(outputOpOperand.value()->get());
         newIndexingMaps.push_back(
             genericOp.getTiedIndexingMap(outputOpOperand.value()));
       }
     }
 
void populateOpPayload(
 
     // 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);
   }