[Global Opt] Raise tensor.extract to input (iree-org#22434)

The PR that added `raiseTensorExtractToInput`
(iree-org#14718) only used the
`linalg.generic` as an intermediate step when trying to convert to a
`tensor` view-like op. However, we should be replace the generic op with
the `tensor.extract`-less generic op when it can't be converted to a
view-like op.

For example, the linalg op in the test case `test_extract_to_transpose`
cannot be converted into a view but should be converted to a simple
transpose.

Signed-off-by: Ian Wood <[email protected]>

Author: IanWood1

compiler/src/iree/compiler/GlobalOptimization/RaiseSpecialOps.cpp

@@ -1016,28 +1016,28 @@ struct RaiseSpecialOpsPass
 
     // First walk the IR and try to raise any slice-like generics to tensor.
     IRRewriter rewriter(context);
-    funcOp->walk([&](linalg::GenericOp op) {
-      linalg::GenericOp linalgOp = op;
+    SmallVector<linalg::GenericOp> genericOps;
+    funcOp->walk([&](linalg::GenericOp op) { genericOps.push_back(op); });
 
+    for (linalg::GenericOp linalgOp : genericOps) {
+      // Try raising to tensor.extract to an input and create an linalg.generic.
       OpBuilder::InsertionGuard guard(rewriter);
-
-      // Try raising to tensor.export and create an intermediate linalg.generic.
-      rewriter.setInsertionPoint(op);
+      rewriter.setInsertionPoint(linalgOp);
       FailureOr<linalg::GenericOp> maybeNewOp =
           raiseTensorExtractToInput(linalgOp, rewriter);
       if (succeeded(maybeNewOp)) {
+        rewriter.replaceOp(linalgOp, *maybeNewOp);
         linalgOp = *maybeNewOp;
       }
 
       // Try raising to a view-like operation. Replace if the op raising was
       // successful.
-      rewriter.setInsertionPoint(op);
       FailureOr<Operation *> maybeRaisedView =
           tryRaiseToView(linalgOp, rewriter);
       if (succeeded(maybeRaisedView)) {
-        rewriter.replaceOp(op, *maybeRaisedView);
+        rewriter.replaceOp(linalgOp, *maybeRaisedView);
       }
-    });
+    }
 
     // Next run a variety of raising patterns.
     {

compiler/src/iree/compiler/GlobalOptimization/test/raise_special_ops.mlir

@@ -293,6 +293,26 @@ util.func public @test_slice_middle(%A : tensor<64x64x64xf32>, %B : tensor<64x64
 
 // -----
 
+#map = affine_map<(d0, d1) -> (d0, d1)>
+util.func public @test_extract_to_transpose(%A : tensor<64x64xf32>, %B : tensor<64x64xf32>) -> tensor<64x64xf32> {
+  %c0 = arith.constant 0 : index
+  %0 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel"]} outs(%B : tensor<64x64xf32>) {
+  ^bb0(%out: f32):
+    %i0 = linalg.index 0 : index
+    %i1 = linalg.index 1 : index
+    %extracted = tensor.extract %A[%i1, %i0] : tensor<64x64xf32>
+    linalg.yield %extracted : f32
+  } -> tensor<64x64xf32>
+  util.return %0 : tensor<64x64xf32>
+}
+
+// CHECK-LABEL: util.func public @test_extract_to_transpose
+//       CHECK:   %[[RESULT:.+]] = linalg.generic
+//  CHECK-SAME:     affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>
+//       CHECK:   util.return %[[RESULT]]
+
+// -----
+
 util.func public @test_trailing_elementwise(%arg0: tensor<180x320x1xf32>) -> tensor<320xf32> {
   %c0 = arith.constant 0 : index
   %c179 = arith.constant 179 : index