[mlir][Linalg] Combine canonicalizers that deal with removing dead/redundant args.

[mravishankar]

Summary

linalg.generic ops have canonicalizers that either remove arguments
not used in the payload, or redundant arguments. Combine these and
enhance the canonicalization to also remove results that have no use.
This is effectively dead code elimination for Linalg ops.

[nirvedhmeshram]

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

16

Do you need this?

21

Same question as above.

705

There is a helper function called hasUniqueDim on line 114 that is not needed with this logic, so you can delete that.

710

is is -> is
maybe you want to mention these two conditions are "or" and not "and"?

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir

94

I believe as part of adding the tests here, you can delete two tests in canonicalize.mlir
https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Linalg/canonicalize.mlir#L346-L392

[okkwon]

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

700

I think this can be put in the function.

705

inversePermutation() asserts on that the indexing map does not have a symbolic expression. We might need an explicit check for that.

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir

108

Q: do we expect a race condition as a desirable outcome?

[mravishankar]

Address comments and disallow dropping of dead outputs when it is needed for loop dim computation.

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

16

Ah thanks. Dont need it anymore. I thought I did.

21

This I do need.

700

Wont it result in a allocation on every invocation?

705

LinalgOp verify checks for no symbols. So it can be an assumed property. In general this computation is how loop bounds are determined based on indexing maps.

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir

94

Should be fine. More tests are not a problem :)

108

Not sure I follow the question. How is there a race condition here?

[okkwon]

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

700

Good point!

705

Thanks for the info!

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir

108

Oh, I think I was confused because %arg0 is used for all outputs. But now it looks fine because only %b0 is used. Sorry for the noise.

[gysit]

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

722

All the uses I found are actually calling getValue immediately to get the AffineMap. I would thus suggest to use the existing getTiedIndexingMap function to avoid growing the interface? If an attribute is needed then AffineMapAttr::get() is pretty compact to get it?

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

168

nit: resultTensorTypses -> resultTensorTypes

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

691

nit: in most places we prefix the operands of the newly created op with new instead of modified. Feel free if you think modified is a better fit here!

721

I think it would be nice if we could avoid the assumption and rely on the structured op interface only?

We could for example use a slightly different algorithm that for every loop dimension counts the number of indexing map result appearances. E.g. if an op has the indexing maps (d0, d1) -> (d0, d1) and (d0, d1) -> (d0) we could get the vector dimCounts = [2, 1] saying d0 appears twice and d1 once. We can then iterate opOperands and and try to subract the indexingMap tied to the current opOperand. If the dimCounts vector contains a zero after the subtraction we cannot delete it. In the example, only the operand tied to (d0, d1) -> (d0) can be removed.

I think if we have two lambdas/functions to add and subtract the result dim counts for a specific indexing map the change may even simplify the pattern a little bit. Admittedly, it also implements some logic that is usually hidden in getShapesToLoopsMap (the real one in the structured op interface).

An alternative could be to have a method on the structured op interface that takes a list of dropped operands and returns if the remaining ones suffice to define the iteration domain E.g.:

genericOp.isIterationDomainDefinedBy(/*droppedOperands*/=SmallVector<OpOperand*>{opOperandToDelete1, opOperandToDelete2});

which could then be used in the verification as well. That could fit well with the current algorithm.

845

I think it should be possible to use rewriter.replaceOp(genericOp, newResults) by setting the dropped values to the old result or nullptr? Or is there some check that prevents this?

[nirvedhmeshram]

LGTM for my comments.

[mravishankar]

Address most comments.

[mravishankar]

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

722

You are right. I dont need this anymore.

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

168

Seems like Gollum came and wrote that signature :)

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

691

Good idea. Smaller names.

721

The first approach will require walking the indexing maps, and maintaining involved data structures to get it current. Seems more invasive and error prone.

The second approach looks feasible. Let me try that.

[mravishankar]

Use interface method to check if operands can be dropped while still being able to compute loop bounds.

[mravishankar]

Rebase.

[mravishankar]

Fix a bug in the canonicalizer and add test that exercises that path.

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

InterfaceMethod<
         assert(opOperand->getOwner() == this->getOperation());
         auto indexingMaps =
           $_op.indexing_maps().template getAsValueRange<AffineMapAttr>();
         return *(indexingMaps.begin() + opOperand->getOperandNumber());
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return the input or output indexing map for `opOperand`.
+      }],
+      /*retTy=*/"AffineMapAttr",
+      /*methodName=*/"getTiedIndexingMapAttr",
+      /*args=*/(ins "OpOperand*":$opOperand),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        assert(opOperand->getOwner() == this->getOperation());
+        Attribute mapAttr = $_op.indexing_maps()[opOperand->getOperandNumber()];
+        return mapAttr.template cast<AffineMapAttr>();
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the indexing map for a `result`.
       }],
       /*retTy=*/"AffineMap",
       /*methodName=*/"getTiedIndexingMapForResult",
       /*args=*/(ins "OpResult":$result),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         assert(result.getOwner() == this->getOperation());

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

let arguments = (ins Variadic<AnyType>:$inputs,
                        AffineMapArrayAttr:$indexing_maps,
                        ArrayAttr:$iterator_types,
                        OptionalAttr<StrAttr>:$doc,
                        OptionalAttr<StrAttr>:$library_call);
   let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
   let regions = (region AnyRegion:$region);
 
   let builders = [
+    OpBuilder<(ins "TypeRange":$resultTensorTypses, "ValueRange":$inputs,
+      "ValueRange":$outputs, "ArrayAttr":$indexingMaps,
+      "ArrayAttr":$iteratorTypes, "StringAttr":$doc,
+      "StringAttr":$libraryCall,
+      "function_ref<void(OpBuilder &, Location, ValueRange)>",
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
     OpBuilder<(ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
       "ValueRange":$outputs, "ArrayRef<AffineMap>":$indexingMaps,
       "ArrayRef<StringRef>":$iteratorTypes, "StringRef":$doc,
       "StringRef":$libraryCall,
       CArg<"function_ref<void(OpBuilder &, Location, ValueRange)>", "nullptr">,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
     OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputBuffers,
       "ArrayRef<AffineMap>":$indexingMaps, "ArrayRef<StringRef>":$iteratorTypes,

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

 //===- LinalgOps.cpp - Implementation of the linalg operations ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the Linalg operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineExprVisitor.h"
+#include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser/Parser.h"
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SetVector.h"
for (OpOperand &operand : op->getOpOperands()) {
     if (castOp && memref::CastOp::canFoldIntoConsumerOp(castOp)) {
       operand.set(castOp.getOperand());
       folded = true;
     }
   }
   return success(folded);
 }
 
-/// Helper function to find if there is atleast one dimension in an AffineMap
-/// testMap that is contained in `testMapLocation` of  `maps` but not in any
-/// other locations
-static bool hasaUniqueDim(ArrayRef<AffineMap> maps, unsigned testMapLocation) {
-  AffineMap testMap = maps[testMapLocation];
-  llvm::SmallDenseSet<unsigned> dimsToCheck;
-  for (auto result : testMap.getResults()) {
-    auto expr = result.dyn_cast<AffineDimExpr>();
-    if (expr != nullptr)
-      dimsToCheck.insert(expr.getPosition());
-  }
-  for (auto It : llvm::enumerate(maps)) {
-    if (It.index() == testMapLocation)
-      continue;
-    auto map = It.value();
-    for (auto result : map.getResults()) {
-      auto expr = result.dyn_cast<AffineDimExpr>();
-      if (expr != nullptr) {
-        dimsToCheck.erase(expr.getPosition());
-      }
-      if (dimsToCheck.empty())
-        return false;
-    }
-  }
-  return true;
-}
-
 //===----------------------------------------------------------------------===//
 // Region builder helper.
 // TODO: Move this to a utility library.
 // The public methods on this class are referenced directly from generated code.
 // Helper build the unary, binary, and type conversion functions defined by the
 // DSL. See mlir-linalg-ods-yaml-gen.cpp for the code that uses this class.
 //
 // Implementations of the math functions must be polymorphic over numeric types,
results
            FoldInsertPadIntoFill>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // GenericOps
 //===----------------------------------------------------------------------===//
 void GenericOp::build(
     OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
-    ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
-    ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
+    ValueRange inputs, ValueRange outputs, ArrayAttr indexingMaps,
+    ArrayAttr iteratorTypes, StringAttr doc, StringAttr libraryCall,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,
     ArrayRef<NamedAttribute> attributes) {
-  build(builder, result, resultTensorTypes, inputs, outputs,
-        builder.getAffineMapArrayAttr(indexingMaps),
-        builder.getStrArrayAttr(iteratorTypes),
-        doc.empty() ? StringAttr() : builder.getStringAttr(doc),
-        libraryCall.empty() ? StringAttr()
-                            : builder.getStringAttr(libraryCall));
+  build(builder, result, resultTensorTypes, inputs, outputs, indexingMaps,
+        iteratorTypes, doc, libraryCall);
   result.addAttributes(attributes);
   if (!bodyBuild)
     return;
 
   SmallVector<Type, 4> blockArgTypes;
   SmallVector<Location, 4> blockArgLocs;
   for (ValueRange container : {inputs, outputs}) {
     for (Value v : container) {
       blockArgTypes.push_back(getElementTypeOrSelf(v));
       blockArgLocs.push_back(v.getLoc());
     }
   }
 
   OpBuilder::InsertionGuard guard(builder);
   auto &region = *result.regions.front();
   Block *bodyBlock =
       builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
   bodyBuild(builder, result.location, bodyBlock->getArguments());
 }
 
 void GenericOp::build(
+    OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
+    ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
+    ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,
+    ArrayRef<NamedAttribute> attributes) {
+  build(builder, result, resultTensorTypes, inputs, outputs,
+        builder.getAffineMapArrayAttr(indexingMaps),
+        builder.getStrArrayAttr(iteratorTypes),
+        doc.empty() ? StringAttr() : builder.getStringAttr(doc),
+        libraryCall.empty() ? StringAttr() : builder.getStringAttr(libraryCall),
+        bodyBuild, attributes);
+}
+
+void GenericOp::build(
     OpBuilder &builder, OperationState &result, ValueRange inputs,
     ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
     ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,
     ArrayRef<NamedAttribute> attributes) {
   build(builder, result, TypeRange{}, inputs, outputs, indexingMaps,
         iteratorTypes, doc, libraryCall, bodyBuild, attributes);
 }

 template <typename GenericOpType>
 static LogicalResult verifyGenericOp(GenericOpType op) {
   return success();
 }
 
 LogicalResult GenericOp::verify() { return verifyGenericOp(*this); }
 
 namespace {
-// Deduplicate redundant args of a linalg generic op.
-// An arg is redundant if it has the same Value and indexing map as another.
-struct DeduplicateGenericOpInputs : public OpRewritePattern<GenericOp> {
+
+struct DeduplicateAndRemoveDeadOperandsAndResults
+    : public OpRewritePattern<GenericOp> {
   using OpRewritePattern<GenericOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
-    // Associate each input to an equivalent "canonical" input that has the same
-    // Value and indexing map.
-    //
-    // In the non-duplicate case, input `i` will have canonical input `i`. But
-    // in the case of duplicated inputs, the canonical input could be some other
-    // input `< i`. That is, a later input will have some earlier input as its
-    // canonical input.
-    llvm::SmallDenseMap<std::pair<Value, AffineMap>, unsigned> canonicalInput;
-    // For later remapping tasks like deduplicating payload block arguments,
-    // having a simple "inputIndex -> canonicalInputIndex" integer mapping is
-    // convenient.
-    SmallVector<unsigned> canonicalInputIndices;
-    for (OpOperand *opOperand : genericOp.getInputOperands()) {
-      AffineMap indexingMap = genericOp.getTiedIndexingMap(opOperand);
-      // STL-like maps have a convenient behavior for our use case here. In the
-      // case of duplicate keys, the insertion is rejected, and the returned
-      // iterator gives access to the value already in the map.
-      auto pair = canonicalInput.insert(
-          {{opOperand->get(), indexingMap}, opOperand->getOperandNumber()});
-      canonicalInputIndices.push_back(pair.first->second);
-    }
-
-    // If there are no duplicate args, then bail out.
-    if (canonicalInput.size() == genericOp.getNumInputs())
-      return failure();
-
-    // The operands for the newly canonicalized op.
-    SmallVector<Value> newInputOperands;
-    for (OpOperand *opOperand : genericOp.getInputOperands())
-      if (canonicalInputIndices[opOperand->getOperandNumber()] ==
-          opOperand->getOperandNumber())
-        newInputOperands.push_back(opOperand->get());
-
-    // Repair the indexing maps by filtering out the ones that have been
-    // eliminated.
-    SmallVector<AffineMap> newIndexingMaps;
-    for (OpOperand *opOperand : genericOp.getInputOperands())
-      if (canonicalInputIndices[opOperand->getOperandNumber()] ==
-          opOperand->getOperandNumber())
-        newIndexingMaps.push_back(genericOp.getTiedIndexingMap(opOperand));
-    for (OpOperand *opOperand : genericOp.getOutputOperands())
-      newIndexingMaps.push_back(genericOp.getTiedIndexingMap(opOperand));
+    // 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.
+    llvm::SmallDenseMap<unsigned, unsigned> origToModifiedPos;
+    unsigned numNewArgs = 0;
+    llvm::SmallDenseSet<unsigned> droppedOutputs;
+
+    // Information needed to build the modified op.
+    SmallVector<Value> modifiedInputOperands, modifiedOutputOperands;
+    SmallVector<AffineMap> modifiedIndexingMaps;
+    SmallVector<Type> modifiedResultTypes;
+
+    // Function to determine if the provided list of indexing maps can compute
+    // the op iteration domain. It takes two lists, one the set of indexing maps
+    // known to be preserved in the new op, and the set of indexing maps from
+    // the original op that might make it into the canonicalized op.
+    SmallVector<AffineMap> concatList;
+    auto getShapesToLoopMap = [&concatList](ArrayRef<AffineMap> list1,
+                                            ArrayRef<AffineMap> list2) -> bool {
+      concatList.assign(list1.begin(), list1.end());
+      concatList.append(list2.begin(), list2.end());
+      return inversePermutation(concatAffineMaps(concatList)) != AffineMap();
+    };
+
+    // Input argument can be dropped if
+    // - it has no uses, or,
+    // - there is a duplicate operand which is accessed using the same
+    //   indexing map.
+    llvm::SmallDenseMap<std::pair<Value, AffineMap>, unsigned> dedupedInputs;
+    auto indexingMaps = genericOp.getIndexingMaps();
+    ArrayRef<AffineMap> unprocessedIndexingMaps(indexingMaps);
+    for (OpOperand *inputOpOperand : genericOp.getInputOperands()) {
+      BlockArgument arg = genericOp.getTiedBlockArgument(inputOpOperand);
+      unsigned argNum = arg.getArgNumber();
+      unprocessedIndexingMaps = unprocessedIndexingMaps.drop_front();
+
+      // Check if operand is dead and if dropping the indexing map makes the
+      // loops to shape computation invalid.
+      // NOTE: This is using the implicit assumption that the first
+      // n indexing maps in the list are for the inputs.
+      if (!genericOp.payloadUsesValueFromOperand(inputOpOperand) &&
+          getShapesToLoopMap(modifiedIndexingMaps, unprocessedIndexingMaps))
+        continue;
+
+      // Check if this operand is a duplicate.
+      AffineMap indexingMap =
+          genericOp.getTiedIndexingMapAttr(inputOpOperand).getValue();
+      auto it = dedupedInputs.find(
+          std::make_pair(inputOpOperand->get(), indexingMap));
+      if (it != dedupedInputs.end()) {
+        origToModifiedPos[argNum] = it->second;
+        continue;
+      }
+
+      // This is a preserved argument.
+      origToModifiedPos[argNum] = numNewArgs++;
+      dedupedInputs[{inputOpOperand->get(), indexingMap}] = argNum;
+      modifiedInputOperands.push_back(inputOpOperand->get());
+      modifiedIndexingMaps.push_back(indexingMap);
+    }
+
+    // If the op doesnt have tensor semantics, keep all the outputs as
+    // preserved.
+    if (!genericOp.hasTensorSemantics()) {
+      for (OpOperand *outputOpOperand : genericOp.getOutputOperands()) {
+        unprocessedIndexingMaps = unprocessedIndexingMaps.drop_front();
+        BlockArgument arg = genericOp.getTiedBlockArgument(outputOpOperand);
+        origToModifiedPos[arg.getArgNumber()] = numNewArgs++;
+        modifiedOutputOperands.push_back(outputOpOperand->get());
+        modifiedIndexingMaps.push_back(
+            genericOp.getTiedIndexingMapAttr(outputOpOperand).getValue());
+      }
+    } 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.
+      for (auto outputOpOperand :
+           llvm::enumerate(genericOp.getOutputOperands())) {
+        unprocessedIndexingMaps = unprocessedIndexingMaps.drop_front();
+        Value result = genericOp.getResult(outputOpOperand.index());
+        BlockArgument arg =
+            genericOp.getTiedBlockArgument(outputOpOperand.value());
+        // NOTE: The loop bounds computation here assumes that the list of
+        // operands is first input list followed by output list.
+        if (result.use_empty() &&
+            !genericOp.payloadUsesValueFromOperand(outputOpOperand.value()) &&
+            getShapesToLoopMap(modifiedIndexingMaps, unprocessedIndexingMaps)) {
+          // Drop the output.
+          droppedOutputs.insert(outputOpOperand.index());
+          continue;
+        }
 
-    // Clone the old op with new operands.
-    SmallVector<Value> outputOperands = genericOp.getOutputOperands();
+        origToModifiedPos[arg.getArgNumber()] = numNewArgs++;
+        modifiedOutputOperands.push_back(outputOpOperand.value()->get());
+        modifiedIndexingMaps.push_back(
+            genericOp.getTiedIndexingMapAttr(outputOpOperand.value())
+                .getValue());
+        modifiedResultTypes.push_back(result.getType());
+      }
+    }
+
+    if (modifiedInputOperands.size() + modifiedOutputOperands.size() ==
+        genericOp.getNumInputsAndOutputs()) {
+      // Nothing to do.
+      return failure();
+    }
+
+    // Create the new op with the body being empty.
+    Location loc = genericOp.getLoc();
     auto newOp = rewriter.create<GenericOp>(
-        genericOp.getLoc(), genericOp->getResultTypes(), newInputOperands,
-        outputOperands, rewriter.getAffineMapArrayAttr(newIndexingMaps),
+        loc, modifiedResultTypes, modifiedInputOperands, modifiedOutputOperands,
+        rewriter.getAffineMapArrayAttr(modifiedIndexingMaps),
         genericOp.iterator_types(), genericOp.docAttr(),
-        genericOp.library_callAttr());
-
+        genericOp.library_callAttr(),
+        [](OpBuilder & /*builder*/, Location /*loc*/, ValueRange /*args*/) {
+          return;
+        });
     // Copy over unknown attributes. They might be load bearing for some flow.
     ArrayRef<StringRef> odsAttrs = genericOp.getAttributeNames();
     for (NamedAttribute kv : genericOp->getAttrs()) {
       if (!llvm::is_contained(odsAttrs, kv.getName().getValue())) {
         newOp->setAttr(kv.getName(), kv.getValue());
       }
     }
 
-    rewriter.inlineRegionBefore(genericOp.region(), newOp.region(),
-                                newOp.region().begin());
-
-    // Repair the payload entry block by RAUW'ing redundant arguments and
-    // erasing them.
-    Block &payload = newOp.region().front();
-    SmallVector<OpOperand *> inputOperands = genericOp.getInputOperands();
-    for (OpOperand *opOperand : llvm::reverse(inputOperands)) {
-      // Iterate in reverse, so that we erase later args first, preventing the
-      // argument list from shifting unexpectedly and invalidating all our
-      // indices.
-      unsigned operandNumber = opOperand->getOperandNumber();
-      if (canonicalInputIndices[operandNumber] == operandNumber)
+    // Merge the body of the original op with the new op.
+    Block *newOpBlock = &newOp.region().front();
+    Block *origOpBlock = &genericOp.region().front();
+    SmallVector<Value> replacements(origOpBlock->getNumArguments(), nullptr);
+    for (auto argNum : llvm::seq<unsigned>(0, origOpBlock->getNumArguments())) {
+      auto it = origToModifiedPos.find(argNum);
+      if (it != origToModifiedPos.end()) {
+        replacements[argNum] = newOpBlock->getArgument(it->second);
+      }
+    }
+    rewriter.mergeBlocks(origOpBlock, newOpBlock, replacements);
+
+    // Drop the unused yield args.
+    Block *block = &newOp.region().front();
+    if (!droppedOutputs.empty()) {
+      OpBuilder::InsertionGuard g(rewriter);
+      SmallVector<Value> newYieldVals;
+      YieldOp origYieldOp = cast<YieldOp>(block->getTerminator());
+      rewriter.setInsertionPoint(origYieldOp);
+      for (auto yieldOpOperands : llvm::enumerate(origYieldOp.values())) {
+        if (!droppedOutputs.count(yieldOpOperands.index())) {
+          newYieldVals.push_back(yieldOpOperands.value());
           continue;
-      payload.getArgument(operandNumber)
-          .replaceAllUsesWith(
-              payload.getArgument(canonicalInputIndices[operandNumber]));
-      payload.eraseArgument(operandNumber);
+        }
+      }
+      rewriter.replaceOpWithNewOp<YieldOp>(origYieldOp, newYieldVals);
     }
 
-    rewriter.replaceOp(genericOp, newOp->getResults());
+    // Replace all live uses of the op.
+    unsigned newResultNum = 0;
+    for (auto result : llvm::enumerate(genericOp.getResults())) {
+      if (droppedOutputs.count(result.index()))
+        continue;
+      result.value().replaceAllUsesWith(newOp.getResult(newResultNum++));
+    }
+    rewriter.eraseOp(genericOp);
     return success();
   }
 };
 
 /// Remove generic operations (on tensors) that are just copying
 /// the values from inputs to the results. Requirements are
 /// 1) All iterator types are parallel
 /// 2) The body contains just a yield operation with the yielded values being
LogicalResult matchAndRewrite(GenericOp genericOp,
     }
 
     if (returnedArgs.size() != genericOp->getNumResults())
       return failure();
     rewriter.replaceOp(genericOp, returnedArgs);
     return success();
   }
 };
-
-/// Drop dead args of a linalg generic op.
-/// An arg is dead if it has zero uses in the op region.
-struct DeadArgsGenericOpInputs : public OpRewritePattern<GenericOp> {
-  using OpRewritePattern<GenericOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(GenericOp genericOp,
-                                PatternRewriter &rewriter) const override {
-    SmallVector<AffineMap> oldIndexingMaps = genericOp.getIndexingMaps();
-    // Maps must be projected permutations.
-    if (llvm::any_of(genericOp.getIndexingMaps(), [](AffineMap map) {
-          return !map.isProjectedPermutation();
-        }))
-      return failure();
-    Block &payload = genericOp.region().front();
-    SmallVector<Value> newInputOperands;
-    SmallVector<AffineMap> newIndexingMaps;
-    bool deadArgFound = false;
-    int inputSize = genericOp.getInputOperands().size();
-    for (int i = inputSize - 1; i >= 0; i--) {
-      OpOperand *opOperand = genericOp.getInputOperand(i);
-      // Iterate in reverse, so that we erase later args first, preventing the
-      // argument list from shifting unexpectedly and invalidating all our
-      // indices.
-      if (payload.getArgument(i).use_empty() &&
-          !hasaUniqueDim(oldIndexingMaps, i)) {
-        payload.eraseArgument(i);
-        deadArgFound = true;
-        // remove this indexing map out of consideration for hasaUniqueDim check
-        oldIndexingMaps.erase(oldIndexingMaps.begin() + i);
-      } else {
-        newInputOperands.insert(newInputOperands.begin(), opOperand->get());
-        newIndexingMaps.insert(newIndexingMaps.begin(),
-                               genericOp.getTiedIndexingMap(opOperand));
-      }
-    }
-    // Bail out if there are no dead args.
-    if (!deadArgFound)
-      return failure();
-    for (OpOperand *opOperand : genericOp.getOutputOperands())
-      newIndexingMaps.push_back(genericOp.getTiedIndexingMap(opOperand));
-    SmallVector<Value> outputOperands = genericOp.getOutputOperands();
-
-    auto newOp = rewriter.create<GenericOp>(
-        genericOp.getLoc(), genericOp->getResultTypes(), newInputOperands,
-        outputOperands, rewriter.getAffineMapArrayAttr(newIndexingMaps),
-        genericOp.iterator_types(), genericOp.docAttr(),
-        genericOp.library_callAttr());
-    // Copy over unknown attributes. They might be load bearing for some flow.
-    ArrayRef<StringRef> odsAttrs = genericOp.getAttributeNames();
-    for (NamedAttribute kv : genericOp->getAttrs()) {
-      if (!llvm::is_contained(odsAttrs, kv.getName().getValue())) {
-        newOp->setAttr(kv.getName(), kv.getValue());
-      }
-    }
-    rewriter.inlineRegionBefore(genericOp.region(), newOp.region(),
-                                newOp.region().begin());
-    rewriter.replaceOp(genericOp, newOp->getResults());
-    return success();
-  }
-};
 } // namespace
 
 void GenericOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
-  results.add<DeduplicateGenericOpInputs, EraseIdentityGenericOp,
-              DeadArgsGenericOpInputs>(context);
+  results
+      .add<DeduplicateAndRemoveDeadOperandsAndResults, EraseIdentityGenericOp>(
+          context);
 }
 
 LogicalResult GenericOp::fold(ArrayRef<Attribute>,
                               SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir

%0 = linalg.generic {indexing_maps = [#map, #map, #map, #map, #map], iterator_types = ["parallel"]}
      ins(%arg0, %arg1, %arg0, %arg1 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>)
     outs(%arg0 : tensor<?xf32>) {
   ^bb0(%arg2: f32, %arg3: f32, %arg4: f32, %arg5: f32, %arg6: f32):
     %1 = "test.elementwise_mappable"(%arg2, %arg3, %arg4, %arg5) : (f32, f32, f32, f32) -> f32
     linalg.yield %1 : f32
   } -> tensor<?xf32>
   return %0 : tensor<?xf32>
 }
+
+// -----
+
+// Drop dead result.
+
+#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
+#map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
+func @drop_dead_results(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+  %0:4 = linalg.generic {
+      indexing_maps = [#map0, #map1, #map2, #map3, #map4],
+      iterator_types = ["parallel", "parallel", "parallel"]}
+      ins(%arg0 : tensor<?x?x?xf32>)
+      outs(%arg0, %arg0, %arg0, %arg0
+          : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) :
+      %1 = arith.addf %b0, %b0: f32
+      linalg.yield %1, %1, %1, %1 : f32, f32, f32, f32
+    } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+  return %0#0, %0#2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>     
+}
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
+//      CHECK: func @drop_dead_results(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?x?xf32>)
+//      CHECK:   %[[GENERIC:.+]]:2 = linalg.generic
+// CHECK-SAME:       indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+// CHECK-SAME:       outs(%[[ARG0]], %[[ARG0]] :
+//      CHECK:   return %[[GENERIC]]#0, %[[GENERIC]]#1
+
+// -----
+
+// Current argmax lowering to `linalg.generic`. Cannot drop the
+// first return even though it isnt used since it has an internal
+// use.
+#map0 = affine_map<(d0) -> (d0)>
+#map1 = affine_map<(d0) -> ()>
+func @argmax_lowering(%arg0 : tensor<?xf32>) -> tensor<i32> {
+  %init0 = linalg.init_tensor [] : tensor<f32>
+  %init1 = linalg.init_tensor [] : tensor<i32>
+  %0:2 = linalg.generic {
+    indexing_maps = [#map0, #map1, #map1],
+    iterator_types = ["reduction"]}
+    ins(%arg0 : tensor<?xf32>)
+    outs(%init0, %init1 : tensor<f32>, tensor<i32>) {
+  ^bb0(%b0: f32, %b1: f32, %b2: i32):
+    %8 = linalg.index 0 : index
+    %9 = arith.index_cast %8 : index to i32
+    %10 = arith.cmpf oge, %b0, %b1 : f32
+    %11 = arith.select %10, %b0, %b1 : f32
+    %12 = arith.cmpf oeq, %b0, %b1 : f32
+    %13 = arith.minsi %9, %b2 : i32
+    %14 = arith.select %10, %9, %b2 : i32
+    %15 = arith.select %12, %13, %14 : i32
+    linalg.yield %11, %15 : f32, i32
+  } -> (tensor<f32>, tensor<i32>)
+  return %0#1 : tensor<i32>
+}
+//      CHECK: func @argmax_lowering(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<?xf32>
+//  CHECK-DAG:   %[[INIT0:.+]] = linalg.init_tensor [] : tensor<f32>
+//  CHECK-DAG:   %[[INIT1:.+]] = linalg.init_tensor [] : tensor<i32>
+//      CHECK:   %[[GENERIC:.+]]:2 = linalg.generic
+// CHECK-SAME:       outs(%[[INIT0]], %[[INIT1]] :
+//      CHECK:   return %[[GENERIC]]#1
+
+// -----
+
+// Do not remove operand needed for loop dim.
+func @loop_dim_operand(%arg0 : tensor<?xf32>) -> tensor<i32> {
+  %cst = arith.constant 0 : i32
+  %init = linalg.init_tensor [] : tensor<i32>
+  %fill = linalg.fill ins(%cst : i32) outs(%init : tensor<i32>) -> tensor<i32>
+  %0 = linalg.generic {
+      indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+      iterator_types = ["reduction"]}
+      ins(%arg0 : tensor<?xf32>) outs(%fill : tensor<i32>) {
+    ^bb0(%b0: f32, %b1: i32):
+      %1 = linalg.index 0 : index
+      %2 = arith.index_cast %1 : index to i32
+      %3 = arith.addi %b1, %2 : i32
+      linalg.yield %3 : i32
+    } -> tensor<i32>
+  return %0 : tensor<i32>
+}
+//      CHECK: func @loop_dim_operand(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<?xf32>
+//      CHECK:   linalg.generic
+// CHECK-SAME:       ins(%[[ARG0]] :
+
+// -----
+
+// Do not remove outs operand needed for loop bound computation.
+func @loop_dim_outs_operand(%arg0 : index) -> tensor<i32> {
+  %cst = arith.constant 0 : i32
+  %init1 = linalg.init_tensor [%arg0] : tensor<?xi32>
+  %init = linalg.init_tensor [] : tensor<i32>
+  %fill = linalg.fill ins(%cst : i32) outs(%init : tensor<i32>) -> tensor<i32>
+  %0:2 = linalg.generic {
+      indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+      iterator_types = ["parallel"]}
+      outs(%init1, %fill : tensor<?xi32>, tensor<i32>) {
+    ^bb0(%b0: i32, %b1: i32):
+      %1 = linalg.index 0 : index
+      %2 = arith.index_cast %1 : index to i32
+      %3 = arith.addi %b1, %2 : i32
+      linalg.yield %2, %3 : i32, i32
+    } -> (tensor<?xi32>, tensor<i32>)
+  return %0#1 : tensor<i32>
+}
+//      CHECK: func @loop_dim_outs_operand(
+// CHECK-SAME:     %[[ARG0:.+]]: index
+//      CHECK:   %[[INIT:.+]] = linalg.init_tensor [%[[ARG0]]]
+//      CHECK:   linalg.generic
+// CHECK-SAME:       outs(%[[INIT]]