[mlir][Linalg] Extend Generic op dedup canonicalization to drop redundant result values.

[mravishankar]

Summary

Depends on D129892

[stellaraccident]

Thanks - makes sense.

[mravishankar]

This is closed by commit https://github.com/llvm/llvm-project/commit/331afccad2133dab5a45bc5ad7f0921b3d22725f

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

private:
   // the canonicalized op.
   // - The preserved output operands list (by reference).
   llvm::SmallDenseMap<unsigned, unsigned>
   deduplicateOutputOperands(GenericOp genericOp,
                             SmallVector<OpOperand *> &droppedOpOperands,
                             SmallVector<Value> &newOutputOperands,
                             SmallVector<AffineMap> &newIndexingMaps) const {
     llvm::SmallDenseMap<unsigned, unsigned> origToNewPos;
+    llvm::SmallDenseMap<std::tuple<Value, AffineMap, Value>, unsigned>
+        dedupedOutpts;
     // If the op doesnt have tensor semantics, keep all the outputs as
     // preserved.
     if (!genericOp.hasTensorSemantics()) {
       for (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 (auto outputOpOperand :
            llvm::enumerate(genericOp.getOutputOperands())) {
         Value result = genericOp.getResult(outputOpOperand.index());
-        if (result.use_empty() &&
-            !genericOp.payloadUsesValueFromOperand(outputOpOperand.value())) {
-          // 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.
+        AffineMap indexingMap =
+            genericOp.getTiedIndexingMap(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();
           }
 
+          // 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;
+          }
+        }
+
         origToNewPos[outputOpOperand.index()] = newOutputOperands.size();
+        dedupedOutpts[key] = newOutputOperands.size();
         newOutputOperands.push_back(outputOpOperand.value()->get());
         newIndexingMaps.push_back(
             genericOp.getTiedIndexingMap(outputOpOperand.value()));
       }
     }
 
     return origToNewPos;
   }

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

 //      CHECK:     %[[T0:.+]] = arith.addi %[[B0]], %[[B1]]
 //      CHECK:     %[[T1:.+]] = arith.addi %[[T0]], %[[B1]]
 //      CHECK:     %[[T2:.+]] = arith.addi %[[T1]], %[[B2]]
 //      CHECK:     %[[T3:.+]] = arith.addi %[[T2]], %[[B3]]
 //      CHECK:     %[[T4:.+]] = arith.addi %[[T3]], %[[B3]]
 //      CHECK:     %[[T5:.+]] = arith.addi %[[T4]], %[[B4]]
 //      CHECK:     linalg.yield %[[T5]]
 //      CHECK:  return %[[RETURN]]
+
+// -----
+
+// Drop redundant results.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @drop_redundant_results(
+    %arg0 : tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
+  %0:2 = linalg.generic {
+    indexing_maps = [#map, #map, #map],
+    iterator_types = ["parallel", "parallel"]}
+    ins(%arg0 : tensor<?x?xf32>)
+    outs(%arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
+      %1 = arith.addf %b0, %b0 : f32
+      linalg.yield %1, %1 : f32, f32
+    } -> (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]]