FuseOp: add use_forall argument that generates scf.forall loops

Author: tkarna

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

@@ -410,13 +410,15 @@ def FuseOp : Op<Transform_Dialect, "structured.fuse",
     (ins TransformHandleTypeInterface:$target,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_interchange,
-         DefaultValuedAttr<BoolAttr, "false">:$apply_cleanup);
+         DefaultValuedAttr<BoolAttr, "false">:$apply_cleanup,
+         DefaultValuedAttr<BoolAttr, "false">:$use_forall);
   let results = (outs TransformHandleTypeInterface:$transformed,
                       Variadic<TransformHandleTypeInterface>:$loops);
 
   let assemblyFormat = [{
     $target ($tile_sizes^)? (`interchange` $tile_interchange^)?
-    (`apply_cleanup` `=` $apply_cleanup^)? attr-dict
+    (`apply_cleanup` `=` $apply_cleanup^)?
+    (`use_forall` `=` $use_forall^)? attr-dict
     `:` functional-type(operands, results)
   }];
   let hasVerifier = 1;

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

@@ -637,6 +637,10 @@ transform::FuseOp::apply(transform::TransformRewriter &rewriter,
 
   scf::SCFTilingOptions tilingOptions;
   tilingOptions.interchangeVector = tileInterchange;
+  bool useForall = getUseForall();
+  tilingOptions.setLoopType(useForall
+                                ? scf::SCFTilingOptions::LoopType::ForallOp
+                                : scf::SCFTilingOptions::LoopType::ForOp);
   SmallVector<OpFoldResult> tileSizesOfr =
       getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);
   tilingOptions = tilingOptions.setTileSizes(tileSizesOfr);
@@ -652,9 +656,11 @@ transform::FuseOp::apply(transform::TransformRewriter &rewriter,
     tileAndFuseOptions.cleanupPatterns = std::move(patterns);
   }
 
+  size_t numLoops =
+      useForall ? 1 : tileSizes.size() - llvm::count(tileSizes, 0);
   LogicalResult result = applyTilingToAll(
-      rewriter, getOperation(), state.getPayloadOps(getTarget()),
-      tileSizes.size() - llvm::count(tileSizes, 0), transformResults,
+      rewriter, getOperation(), state.getPayloadOps(getTarget()), numLoops,
+      transformResults,
       [&](TilingInterface tilingInterfaceOp)
           -> FailureOr<scf::SCFTileAndFuseResult> {
         return tileConsumerAndFuseProducersUsingSCF(rewriter, tilingInterfaceOp,
@@ -676,7 +682,8 @@ LogicalResult transform::FuseOp::verify() {
 
   SmallVector<int64_t> sizes =
       extractFromIntegerArrayAttr<int64_t>(getTileSizes());
-  size_t numExpectedLoops = sizes.size() - llvm::count(sizes, 0);
+  size_t numExpectedLoops =
+      getUseForall() ? 1 : sizes.size() - llvm::count(sizes, 0);
   if (numExpectedLoops != getNumResults() - 1)
     return emitOpError() << "expects " << numExpectedLoops << " loop results";
 

mlir/python/mlir/dialects/transform/structured.py

@@ -147,6 +147,7 @@ def __init__(
         tile_sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
         tile_interchange: OptionalIntList = None,
         apply_cleanup: Optional[bool] = False,
+        use_forall: Optional[bool] = False,
         loc=None,
         ip=None,
     ):
@@ -160,6 +161,7 @@ def __init__(
         tile_sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
         tile_interchange: OptionalIntList = None,
         apply_cleanup: Optional[bool] = False,
+        use_forall: Optional[bool] = False,
         loc=None,
         ip=None,
     ):
@@ -173,14 +175,15 @@ def __init__(
         tile_sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
         tile_interchange: OptionalIntList = None,
         apply_cleanup: Optional[bool] = False,
+        use_forall: Optional[bool] = False,
         loc=None,
         ip=None,
     ):
         tile_sizes = tile_sizes if tile_sizes else []
         tile_interchange = tile_interchange if tile_interchange else []
         _, tile_sizes, _ = _dispatch_dynamic_index_list(tile_sizes)
         _, tile_interchange, _ = _dispatch_dynamic_index_list(tile_interchange)
-        num_loops = sum(0 if v == 0 else 1 for v in tile_sizes)
+        num_loops = 1 if use_forall else sum(0 if v == 0 else 1 for v in tile_sizes)
 
         if isinstance(loop_types_or_target, (Operation, Value, OpView)):
             loop_types = [transform.AnyOpType.get()] * num_loops
@@ -200,6 +203,7 @@ def __init__(
             tile_sizes=tile_sizes,
             tile_interchange=tile_interchange,
             apply_cleanup=apply_cleanup,
+            use_forall=use_forall,
             loc=loc,
             ip=ip,
         )

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

@@ -57,6 +57,31 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// CHECK-LABEL: func.func @fuse_unary_forall
+func.func @fuse_unary_forall(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+
+  //     CHECK: %[[RES:.*]] = scf.forall
+  //     CHECK:       linalg.exp
+  //     CHECK:       linalg.add
+  //     CHECK: return %[[RES]]
+  %0 = linalg.exp ins(%arg0 : tensor<?x?xf32>)
+                             outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
+  %1 = linalg.add ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)
+                             outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.add"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1, %loop = transform.structured.fuse %0 [32, 32] use_forall = true
+      : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+      transform.yield
+  }
+}
+
+// -----
+
 // CHECK-LABEL: func.func @interchange_reduction
 //  CHECK-SAME: (%[[INPUT:.+]]: tensor<12x7x25xf32>)
 func.func @interchange_reduction(%input: tensor<12x7x25xf32>) -> tensor<12x25xf32> {

mlir/test/python/dialects/transform_structured_ext.py

@@ -114,6 +114,19 @@ def testFuseOpCompact(target):
     # CHECK-SAME: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
 
 
+@run
+@create_sequence
+def testFuseOpCompactForall(target):
+    structured.FuseOp(
+        target, tile_sizes=[4, 8], apply_cleanup=True, use_forall=True,
+    )
+    # CHECK-LABEL: TEST: testFuseOpCompact
+    # CHECK: transform.sequence
+    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.fuse %{{.*}}[4, 8]
+    # CHECK-SAME: apply_cleanup = true use_forall = true
+    # CHECK-SAME: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+
+
 @run
 @create_sequence
 def testFuseOpNoArg(target):