[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

1035

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

876

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

878

This can be an assert I think

882

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

885–895

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

885–895

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

885–895

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

885–895

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

861

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;

880

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

885

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?

1232

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

860

this should be a static function

1269

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

1269

I think it's to in

1269

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

[ThomasRaoux]

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

1269

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

1268

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

1273

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

860

done, thanks!

861

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);
 }
 
+static bool isResultValueDead(linalg::GenericOp genericOp, OpResult result) {
+  if (!result.use_empty())
+    return false;
+  // If out operand not used in payload, we can drop it.
+  OpOperand *outputOpOperand =
+      genericOp.getOutputOperand(result.getResultNumber());
+  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.
+  // - User of arg is yield.
+  // - outArg data is not being used by other outArgs.
+
+  // Check block arg and cycle from out operand has a single use.
+  BlockArgument outputArg =
+      genericOp.getRegionOutputArgs()[result.getResultNumber()];
+  if (!outputArg.hasOneUse())
+    return false;
+  Operation *argUserOp = *outputArg.user_begin();
+
+  // Check argUser has no other use.
+  if (!argUserOp->use_empty())
+    return false;
+
+  // Check that argUser is a yield.
+  auto yieldOp = dyn_cast<linalg::YieldOp>(argUserOp);
+  if (!yieldOp)
+    return false;
+
+  // Check outArg data is not being used by other outArgs.
+  if (yieldOp.getOperand(result.getResultNumber()) != outputArg)
+    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,
+      return origToNewPos;
+    }
+      OpResult result = genericOp.getTiedOpResult(outputOpOperand.value());
+      auto key = std::make_tuple(outputOpOperand.value()->get(), indexingMap,
                                  yieldOp->getOperand(outputOpOperand.index()));
-
-        // 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();
-          }
+      assert(genericOp.getNumOutputs() >= outputOpOperand.index() &&
+             "Output op idx greater than number of outputs.");
+      if (isResultValueDead(genericOp, 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())
               continue;
             OpOperand *newOperand = newOperands[it->second];
             replacements[origOperand.value()->getOperandNumber()] =
                 newOpBlock->getArgument(newOperand->getOperandNumber());
           }
         };
 
     OpOperandVector origInputOperands = genericOp.getInputOperands();
     OpOperandVector newInputOperands = newOp.getInputOperands();
     updateReplacements(origInputOperands, newInputOperands, origInsToNewInsPos);
 
     OpOperandVector origOutputOperands = genericOp.getOutputOperands();
     OpOperandVector newOutputOperands = newOp.getOutputOperands();
     updateReplacements(origOutputOperands, newOutputOperands,
                        origOutsToNewOutsPos);
 
-    rewriter.mergeBlocks(origOpBlock, newOpBlock, replacements);
-
     // Drop the unused yield args.
     if (newOp.getNumOutputs() != genericOp.getNumOutputs()) {
       OpBuilder::InsertionGuard g(rewriter);
-      YieldOp origYieldOp = cast<YieldOp>(newOpBlock->getTerminator());
+      YieldOp origYieldOp = cast<YieldOp>(origOpBlock->getTerminator());
       rewriter.setInsertionPoint(origYieldOp);
 
       SmallVector<Value> newYieldVals(newOp.getNumOutputs(), nullptr);
       for (const auto &yieldOpOperands :
            llvm::enumerate(origYieldOp.getValues())) {
         auto it = origOutsToNewOutsPos.find(yieldOpOperands.index());
         if (it == origOutsToNewOutsPos.end())
           continue;
         newYieldVals[it->second] = yieldOpOperands.value();
       }
       rewriter.replaceOpWithNewOp<YieldOp>(origYieldOp, newYieldVals);
     }
+
+    rewriter.mergeBlocks(origOpBlock, newOpBlock, replacements);
   }
 };
 
 /// 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
 ///    the arguments corresponding to the operands.
LogicalResult matchAndRewrite(GenericOp genericOp,
     }
 
     if (returnedArgs.size() != genericOp->getNumResults())
       return failure();
     rewriter.replaceOp(genericOp, returnedArgs);
     return success();
   }
 };
+
+/// Remove unused cycles.
+/// We can remove unused cycle within a payload of generic region
+/// 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 and it is in yield.
+struct RemoveUnusedCycleInGenericOp : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+
+    // If the op doesnt have tensor semantics, preserve the outputs as is.
+    if (!genericOp.hasTensorSemantics())
+      return failure();
+
+    bool hasRemovedCycles = false;
+    // Iterate over output operands and remove any unused cycles.
+    for (const auto &outputOpOperand :
+         llvm::enumerate(genericOp.getOutputOperands())) {
+
+      // Check that result from out operand is dead.
+      Value result = genericOp.getResult(outputOpOperand.index());
+      if (!result.use_empty())
+        continue;
+
+      // Check that outputArg has one use in cycle.
+      BlockArgument outputArg =
+          genericOp.getRegionOutputArgs()[outputOpOperand.index()];
+      if (!outputArg.hasOneUse())
+        continue;
+
+      // Check cycle has at most one use.
+      Operation *cycleOp = *outputArg.user_begin();
+      if (!cycleOp->hasOneUse())
+        continue;
+
+      // Check that the cycleUser is a yield.
+      Operation *cycleUserOp = *cycleOp->user_begin();
+      if (!isa<linalg::YieldOp>(cycleUserOp))
+        continue;
+
+      // Check that argIndex matches yieldIndex, else data is being used.
+      if (cycleUserOp->getOperand(outputOpOperand.index()) !=
+          cycleOp->getResult(0))
+        continue;
+
+      // Directly replace the cycle with the blockArg such that
+      // Deduplicate pattern can eliminate it along with unused yield.
+      rewriter.replaceOp(cycleOp, outputArg);
+      rewriter.updateRootInPlace(genericOp, [] {});
+      hasRemovedCycles = true;
+    }
+
+    if (hasRemovedCycles) {
+      return success();
+    }
+
+    return failure();
+  }
+};
 } // namespace
 
 void GenericOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
-  results
-      .add<DeduplicateAndRemoveDeadOperandsAndResults, EraseIdentityGenericOp>(
-          context);
+  results.add<DeduplicateAndRemoveDeadOperandsAndResults,
+              EraseIdentityGenericOp, RemoveUnusedCycleInGenericOp>(context);
 }
 
 LogicalResult GenericOp::fold(ArrayRef<Attribute>,
                               SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//

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

%0:2 = linalg.generic {
     } -> (tensor<?x?xf32>, tensor<?x?xf32>)
   return %0#0, %0#1 : tensor<?x?xf32>, tensor<?x?xf32>
 }
 //      CHECK: func @drop_redundant_results
 // CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?xf32>
 //      CHECK:   %[[GENERIC:.+]] = linalg.generic
 // CHECK-SAME:       outs(%[[ARG0]] :
 //      CHECK:   return %[[GENERIC]]
+
+// -----
+
+// Drop dead result with different tensors.
+
+#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.func @drop_dead_results_with_different_tensors(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
+  %c1 = arith.constant 1 : index
+  %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
+  %c2 = arith.constant 2 : index
+  %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
+  %init0 = tensor.empty(%d0, %d1, %d2) : 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, %init0, %init0
+          : 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) :
+      linalg.yield %b0, %b0, %b3, %b4 : f32, f32, f32, f32
+    } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+  return %0#0, %0#1 : 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, d2, d0)>
+//      CHECK: func @drop_dead_results_with_different_tensors(
+// 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
+
+// -----
+
+// Drop dead result with unused cycles.
+
+#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.func @drop_dead_results_with_unused_cycles(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
+  %c1 = arith.constant 1 : index
+  %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
+  %c2 = arith.constant 2 : index
+  %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
+  %init0 = tensor.empty(%d0, %d1, %d2) : 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, %init0, %init0
+          : 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
+      %2 = arith.addf %b0, %b3: f32
+      %3 = arith.addf %b0, %b4: f32
+      linalg.yield %1, %1, %2, %3 : f32, f32, f32, f32
+    } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+  return %0#0, %0#1 : 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, d2, d0)>
+//      CHECK: func @drop_dead_results_with_unused_cycles(
+// 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
+
+// -----
+
+// Drop only the results not used by others.
+
+#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)>
+func.func @drop_only_the_results_not_used_by_others(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
+  %c1 = arith.constant 1 : index
+  %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
+  %c2 = arith.constant 2 : index
+  %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
+  %init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32>
+  %0:3 = linalg.generic {
+      indexing_maps = [#map0, #map1, #map2, #map3],
+      iterator_types = ["parallel", "parallel", "parallel"]}
+      ins(%arg0 : tensor<?x?x?xf32>)
+      outs(%arg0, %init0, %init0
+          : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32) :
+      linalg.yield %b2, %b1, %b3 : f32, f32, f32
+    } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+  return %0#0 : tensor<?x?x?xf32>
+}
+
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
+//      CHECK: func @drop_only_the_results_not_used_by_others(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?x?xf32>)
+//      CHECK:   %[[INIT:.+]] = tensor.empty
+//      CHECK:   %[[GENERIC:.+]]:2 = linalg.generic
+// CHECK-SAME:       indexing_maps = [#[[MAP1]], #[[MAP2]]]
+// CHECK-SAME:       outs(%[[ARG0]], %[[INIT]] :
+//      CHECK:   return %[[GENERIC]]#0
+
+// -----
+
+// Drop only the cycles not used by others.
+
+#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)>
+func.func @drop_only_the_cycles_not_used_by_others(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
+  %c1 = arith.constant 1 : index
+  %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
+  %c2 = arith.constant 2 : index
+  %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
+  %init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32>
+  %0:3 = linalg.generic {
+      indexing_maps = [#map0, #map1, #map2, #map3],
+      iterator_types = ["parallel", "parallel", "parallel"]}
+      ins(%arg0 : tensor<?x?x?xf32>)
+      outs(%arg0, %init0, %init0
+          : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32) :
+      %1 = arith.addf %b1, %b2: f32
+      %2 = arith.addf %b1, %b3 : f32
+      linalg.yield %1, %b1, %2 : f32, f32, f32
+    } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+  return %0#0 : tensor<?x?x?xf32>
+}
+
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
+//      CHECK: func @drop_only_the_cycles_not_used_by_others(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?x?xf32>)
+//      CHECK:   %[[INIT:.+]] = tensor.empty
+//      CHECK:   %[[GENERIC:.+]]:2 = linalg.generic
+// CHECK-SAME:       indexing_maps = [#[[MAP1]], #[[MAP2]]]
+// CHECK-SAME:       outs(%[[ARG0]], %[[INIT]] :
+//      CHECK:   return %[[GENERIC]]#0