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
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();
}
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);
}
// 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)
}
}
// 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)
}
}
// 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)
}
}
// 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)
}
}
// 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)
}
}
%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)
}
}
// 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)
}
}
// 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)
}
}
// 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">)
}
}