[mlir] Return new scf.forall handle in fuse_into_containing_op

[harsh]

Summary

Since the scf.forall is now consumed by the fuse into
containing op, we need to return a handle to the new scf.forall.
This patch does that and also ensures that the new bbArg
added to the scf.forall is used in its body.

[ThomasRaoux]

Looks good!

[harsh]

In D151418#4372424, @ThomasRaoux wrote:

Looks good!

Thanks!

[ftynse]

This (and the previous change) breaks transform composability in a quite annoying way. Since the loop handle is now consumed, this invalidates all handles pointing inside the loop. So if we wanted to tile an operation the second time after fusion, we can't anymore (without searching for it again) because we have destroyed the handle to it.

[ftynse]

This also did not update the documentation that keeps saying the operand is only read.

[ftynse]

I will be undoing the "consume" part in https://reviews.llvm.org/D151555, but please update the op documentation to reflect changes made here and in the previous commit.

[harsh]

In D151418#4375866, @ftynse wrote:

I will be undoing the "consume" part in https://reviews.llvm.org/D151555, but please update the op documentation to reflect changes made here and in the previous commit.

Thanks! I am out of the office this week but will get to this next Friday.

[harsh]

Here is the PR for the updated documentation - https://reviews.llvm.org/D152044. Please take a look when you get a chance. Thanks!

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let description = [{
     op are rejected by this operation.
 
     This operation consumes the producer handle.
     This operation only reads the containing op handle.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$producer_op,
                        TransformHandleTypeInterface:$containing_op);
-  let results = (outs TransformHandleTypeInterface:$fused_op);
+  let results = (outs TransformHandleTypeInterface:$fused_op,
+                      TransformHandleTypeInterface:$new_containing_op);
   let assemblyFormat = "$producer_op `into` $containing_op attr-dict "
                        " `:` functional-type(operands, results)";
 
   let builders = [
     OpBuilder<(ins "Value":$producerOp, "Value":$containingOp)>
   ];
 }
 

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 // FuseIntoContainingOp
 //===----------------------------------------------------------------------===//
 
 void transform::FuseIntoContainingOp::build(OpBuilder &builder,
                                             OperationState &result,
                                             Value producerOp,
                                             Value containingOp) {
   result.addOperands({producerOp, containingOp});
-  result.addTypes(transform::AnyOpType::get(builder.getContext()));
+  auto resultType = transform::AnyOpType::get(builder.getContext());
+  result.addTypes({resultType, resultType});
 }
 
 /// Add new operands to the forall op for users of the producerOp
 /// that are dominated by the containing scf.forall op.
 static Operation *replaceForAllWithNewSignature(
     RewriterBase &rewriter, Diagnostic &diag, Operation *producerOp,
     Operation *containingOp, TilingResult &tileAndFuseResult,
     int64_t resultNumber, SmallVector<OpFoldResult> &offsets,
static Operation *replaceForAllWithNewSignature(
   // Create new scf.forall op
   auto newforallOp = rewriter.create<scf::ForallOp>(
       loc, forallOp.getMixedLowerBound(), forallOp.getMixedUpperBound(),
       forallOp.getMixedStep(), newOuts, forallOp.getMapping());
   rewriter.eraseBlock(newforallOp.getBody());
   newforallOp.getRegion().takeBody(forallOp.getRegion());
 
   // Add additional block argument for new value being returned
+  // and replaces all uses of the new output with corresponding bbArg
+  // inside the scf.forall to enable fusion into this new scf.forall.
   newforallOp.getBody()->addArgument(newOuts.back().getType(),
                                      newOuts.back().getLoc());
+  auto bbArgs = newforallOp.getBody()->getArguments();
+  rewriter.replaceUsesWithIf(newOuts.back(), bbArgs.back(),
+                             [&](OpOperand &use) {
+                               Operation *op = use.getOwner();
+                               return newforallOp->isProperAncestor(op);
+                             });
 
   // Fix terminator
   scf::InParallelOp terminatorOp = newforallOp.getTerminator();
   SmallVector<Operation *> yieldingOps = llvm::to_vector<4>(llvm::map_range(
       terminatorOp.getYieldingOps(), [](Operation &op) { return &op; }));
   Operation *firstYieldOp = yieldingOps.front();
   rewriter.setInsertionPoint(firstYieldOp);
   Value src = tileAndFuseResult.tiledValues[0];
if (cloned) {
       LLVM_DEBUG(DBGS() << "\nFused an use by cloning\n" << *containingOp);
       fusedOps.push_back(cloned);
       continue;
     }
     return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
   }
 
   results.set(cast<OpResult>(getFusedOp()), fusedOps);
+  results.set(cast<OpResult>(getNewContainingOp()), {containingOp});
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::FuseIntoContainingOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getProducerOp(), effects);
-  onlyReadsHandle(getContainingOp(), effects);
-  producesHandle(getFusedOp(), effects);
+  consumesHandle(getContainingOp(), effects);
+  producesHandle(getResults(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // GeneralizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure

mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir

module {
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.fill">
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">
 
     // linalg.fill is tileable. The op is tiled and fused.
     transform.structured.fuse_into_containing_op %0 into %1
-      : (!transform.op<"linalg.fill">, !transform.op<"scf.forall">) -> !transform.any_op
+      : (!transform.op<"linalg.fill">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)
   }
 }
 
 // -----
 
 #map0 = affine_map<()[s0] -> (64 ceildiv s0)>
 #map1 = affine_map<(d0)[s0] -> (d0 * s0)>
 #map2 = affine_map<(d0)[s0] -> (-(d0 * s0) + 64, s0)>
module {
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["tensor.empty"]} in %arg1 : (!transform.any_op) -> !transform.op<"tensor.empty">
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">
 
     // tensor.empty is not tileable. The op is cloned and fused.
     transform.structured.fuse_into_containing_op %0 into %1
-      : (!transform.op<"tensor.empty">, !transform.op<"scf.forall">) -> !transform.any_op
+      : (!transform.op<"tensor.empty">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)
   }
 }
 
 // -----
 
 module {
   func.func @foo(%0: tensor<f32>) -> tensor<f32> {
     return %0: tensor<f32>
module {
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.fill">
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">
 
     // linalg.fill is tileable. The op is tiled and fused.
     transform.structured.fuse_into_containing_op %0 into %1
-      : (!transform.op<"linalg.fill">, !transform.op<"scf.forall">) -> !transform.any_op
+      : (!transform.op<"linalg.fill">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)
   }
 }
 
 // -----
 
 #map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
 #map1 = affine_map<(d0)[s0] -> (d0 * s0)>
 #map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
module {
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.any_op
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op
 
     // linalg.fill is tileable. The op is tiled and fused.
     transform.structured.fuse_into_containing_op %0 into %1
-      : (!transform.any_op, !transform.any_op) -> !transform.any_op
+      : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
   }
 }
 
 // -----
 
 #map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
 #map1 = affine_map<(d0)[s0] -> (d0 * s0)>
 #map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
module {
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">
 
     // linalg.generic is tileable. The op is tiled and fused.
     transform.structured.fuse_into_containing_op %0 into %1
-      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> !transform.any_op
+      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)
   }
 }
 
 // -----
 
 module {
   // CHECK-LABEL: func.func @fuse_repeated
   func.func @fuse_repeated(%fill: tensor<2xf32>, %output: tensor<2xf32>) -> tensor<2xf32> {
module {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.any_op
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op
 
     // Create a new handle that points to `linalg.fill` twice.
     %2 = transform.merge_handles %0, %0 : !transform.any_op
 
     // It shouldn't be a problem to fuse this handle.
-    transform.structured.fuse_into_containing_op %2 into %1 : (!transform.any_op, !transform.any_op) -> !transform.any_op
+    transform.structured.fuse_into_containing_op %2 into %1 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
   }
 }
 
 // -----
 
 #map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
 #map1 = affine_map<(d0)[s0] -> (d0 * s0)>
 #map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
module {
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">
 
     // linalg.generic is tileable. The op is tiled and fused.
     transform.structured.fuse_into_containing_op %0 into %1
-      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> !transform.any_op
+      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)
   }
 }
 
 // -----
 
 #map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
 #map1 = affine_map<(d0)[s0] -> (d0 * s0)>
 #map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
module {
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">
 
     // linalg.generic is tileable. The op is tiled and fused.
     transform.structured.fuse_into_containing_op %0 into %1
-      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> !transform.any_op
+      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)
   }
 }
 
 // -----
 
 #map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
 #map1 = affine_map<(d0)[s0] -> (d0 * s0)>
 #map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
module {
 
   transform.sequence failures(propagate) {
   ^bb1(%arg1: !transform.any_op):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">
     %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">
 
     // linalg.generic is tileable. The op is tiled and fused.
     transform.structured.fuse_into_containing_op %0 into %1
-      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> !transform.any_op
+      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)
+  }
+}
+
+// -----
+
+#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
+#map1 = affine_map<(d0)[s0] -> (d0 * s0)>
+#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
+#map3 = affine_map<(d0) -> (d0)>
+
+module {
+  // CHECK-LABEL: func.func @fuse_tileable_using_new_handle
+  //  CHECK-SAME:   %[[CHUNK_SIZE:[0-9a-z]+]]: index
+  //  CHECK-SAME:   %[[IN:[0-9a-z]+]]: tensor<?xf32>
+  //  CHECK-SAME:   %[[OUT_1:[0-9a-z]+]]: tensor<?xf32>
+  //  CHECK-SAME:   %[[OUT_2:[0-9a-z]+]]: tensor<?xf32>
+  //  CHECK-SAME:   %[[OUT_3:[0-9a-z]+]]: tensor<?xf32>
+  func.func @fuse_tileable_using_new_handle(%idx: index, %in: tensor<?xf32>, %out_1: tensor<?xf32>, %out_2: tensor<?xf32>, %out_3: tensor<?xf32>)
+    -> (tensor<?xf32>, tensor<?xf32>) {
+    %cst = arith.constant 4.200000e+01 : f32
+    %c0 = arith.constant 0 : index
+
+    %0 = linalg.generic {
+      indexing_maps = [#map3, #map3], iterator_types = ["parallel"]
+      } ins(%in : tensor<?xf32>) outs(%out_1 : tensor<?xf32>) {
+        ^bb0(%a: f32, %b: f32):
+          %d = arith.addf %a, %b : f32
+          linalg.yield %d : f32
+        } -> tensor<?xf32>
+
+    %1 = linalg.generic {
+      indexing_maps = [#map3, #map3], iterator_types = ["parallel"]
+      } ins(%0 : tensor<?xf32>) outs(%out_1 : tensor<?xf32>) {
+        ^bb0(%a: f32, %b: f32):
+          %d = arith.mulf %a, %b : f32
+          linalg.yield %d : f32
+        } -> tensor<?xf32>
+    %d0 = tensor.dim %out_1, %c0 : tensor<?xf32>
+
+    %2 = affine.apply #map0()[%d0, %idx]
+
+    // CHECK: %[[R0:.*]]:2 = scf.forall (%[[ARG5:.*]]) in (%{{.*}}) shared_outs(%[[ARG6:.*]] = %[[OUT_2]], %[[ARG7:.*]] = %[[OUT_1]])
+    // CHECK-SAME: -> (tensor<?xf32>, tensor<?xf32>) {
+    %3 = scf.forall (%i) in (%2) shared_outs(%o = %out_2) -> (tensor<?xf32>) {
+      // CHECK: %[[I0:.*]] = affine.apply {{.*}}
+      %4 = affine.apply #map1(%i)[%idx]
+      // CHECK: %[[I1:.*]] = affine.min {{.*}}
+      %5 = affine.min #map2(%i)[%d0, %idx]
+      %6 = tensor.extract_slice %o[%4] [%5] [1] : tensor<?xf32> to tensor<?xf32>
+
+      // CHECK: %[[T1:.*]] = linalg.generic {{.*}}
+      // CHECK: %[[T2:.*]] = linalg.generic {{.*}}
+      %7 = tensor.extract_slice %1[%4] [%5] [1] : tensor<?xf32> to tensor<?xf32>
+
+      %8 = linalg.elemwise_unary ins(%7 : tensor<?xf32>) outs(%6 : tensor<?xf32>) -> tensor<?xf32>
+      scf.forall.in_parallel {
+        // CHECK: tensor.parallel_insert_slice %[[T2]] into %[[ARG7]][%[[I0]]] [%[[I1]]] [1] : tensor<?xf32> into tensor<?xf32>
+        tensor.parallel_insert_slice %8 into %o[%2] [%5] [1] : tensor<?xf32> into tensor<?xf32>
+      }
+    }
+    // CHECK: return %[[R0]]#0, %[[R0]]#1
+    func.return %3, %1 : tensor<?xf32>, tensor<?xf32>
+    // CHECK: }
+  }
+
+  transform.sequence failures(propagate) {
+  ^bb1(%arg1: !transform.any_op):
+    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">
+    %add, %reduce = transform.split_handle %0 : (!transform.op<"linalg.generic">) -> (!transform.op<"linalg.generic">, !transform.op<"linalg.generic">)
+    %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">
+
+    %fused_ops, %new_forall = transform.structured.fuse_into_containing_op %reduce into %1
+      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.op<"scf.forall">)
+    %fused_ops_2, %new_forall_2 = transform.structured.fuse_into_containing_op %add into %new_forall
+      : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.op<"scf.forall">)
   }
 }

mlir/test/Dialect/Linalg/transform-tile-and-fuse.mlir

^bb1(%arg1: !transform.any_op):
     %producers = transform.structured.match attributes{"__producer__"} in %arg1 : (!transform.any_op) -> !transform.any_op
 
     // Tile the root.
     %forall_op, %tiled_op = transform.structured.tile_to_forall_op %root num_threads [10, 20]
          : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
     // Fuse all producers.
     transform.structured.fuse_into_containing_op %producers into %forall_op
-      : (!transform.any_op, !transform.any_op) -> !transform.any_op
+      : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
   }
 }
 
 // -----
 
 // Inverse the order of the payload ops passed to the tile_to_forall_op
 // op. Fusion should still work.
 
^bb1(%arg1: !transform.any_op):
     %reversed_producers = transform.test_reverse_payload_ops %producers : (!transform.any_op) -> !transform.any_op
 
     // Tile the root.
     %forall_op, %tiled_op = transform.structured.tile_to_forall_op %root num_threads [10, 20]
          : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
     // Fuse all producers.
     transform.structured.fuse_into_containing_op %reversed_producers into %forall_op
-      : (!transform.any_op, !transform.any_op) -> !transform.any_op
+      : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
   }
 }