[flang-commits] [mlir] [flang] [mlir][TilingInterface] Use `LoopLikeOpInterface` in tiling using SCF to unify tiling with `scf.for` and `scf.forall`. (PR #77874)

Mehdi Amini via flang-commits flang-commits at lists.llvm.org
Thu Jan 18 00:47:36 PST 2024


================
@@ -218,6 +221,54 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
         return ::mlir::failure();
       }]
     >,
+    InterfaceMethod<[{
+        Append the specified additional "init" operands: replace this loop with
+        a new loop that has the additional init operands. The loop body of
+        this loop is moved over to the new loop.
+
+        This method is similar to `replaceWithAdditionalYields` but instead of
+        yielding a value from within the loop, it allows each loop construct
+        implementing this method to handle the result of each iteration
+        appropriately. This allows for unified handling of operations
+        like `scf.forall` which don't yield a value from the loop, but instead
+        the terminator specifies where to insert the tile computed by the body of
+        the loop. For example,
+        
+        ```mlir
+        %0 = scf.forall ... shared_outs(%arg0 = %arg1) {
+          ...
+          %tiled_value = ...
+          scf.forall.in_parallel {
+            tensor.parallel_insert_slice %tiled_value into %arg0[%o1, %o2]...
+          }
+        }
+        ```
+
+        For an `scf.for` the same computation would be represented as
+        ```mlir
+        %0 = scf.for ... iter_args(%arg0 = %arg1) {
+          ...
+          %tiled_value = ...
+          %insert = tensor.insert_slice %tiled_value into %arg0[%o1, %o2]...
+          scf.yield %insert
+        }
+        ```
+
+        So for the caller, the tiled value (`%tiled_value`) and the offsets
+        `(%o1, %o2)` and sizes (not shown) are generated the same way, but
+        the implementation method for the different loop constructs handles
+        the difference in representation.
+      }],
+      /*retTy=*/"::mlir::FailureOr<::mlir::LoopLikeOpInterface>",
+      /*methodName=*/"yieldTiledValuesAndReplace",
+      /*args=*/(ins "::mlir::RewriterBase &":$rewriter,
+                    "::mlir::ValueRange":$newInitOperands,
+                    "const ::mlir::YieldTiledValuesFn &":$yieldTiledValuesFn),
----------------
joker-eph wrote:

```suggestion
                    "::mlir::YieldTiledValuesFn":$yieldTiledValuesFn),
```

`function_ref` is value-based AFAIK (the API using `NewYieldValuesFn` should also be updated)
```suggestion
                    "const ::mlir::YieldTiledValuesFn &":$yieldTiledValuesFn),
```
```suggestion
                    "const ::mlir::YieldTiledValuesFn &":$yieldTiledValuesFn),
```
```suggestion
                    "const ::mlir::YieldTiledValuesFn &":$yieldTiledValuesFn),
```

https://github.com/llvm/llvm-project/pull/77874


More information about the flang-commits mailing list