[Mlir-commits] [mlir] [MLIR][SCF] Fold dim ops of iter_args to respective init_args (PR #109973)

[llvmlistbot at llvm.org]

================
@@ -1478,6 +1478,45 @@ struct DimOfForallOp : public OpRewritePattern<tensor::DimOp> {
   }
 };
 
+/// Fold dim ops of iter_args to dim ops of their respective init args. E.g.:
+///
+/// ```
+/// %0 = ... : tensor<?x?xf32>
+/// scf.forall ... shared_outs(%arg0 = %0) -> (tensor<?x?xf32>) {
+///   %1 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+///   ...
+/// }
+/// ```
+///
+/// is folded to:
+///
+/// ```
+/// %0 = ... : tensor<?x?xf32>
+/// scf.forall ... shared_outs(%arg0 = %0) -> (tensor<?x?xf32>) {
+///   %1 = tensor.dim %0, %c0 : tensor<?x?xf32>
+///   ...
+/// }
+/// ```
+struct DimOfForallIterArg : public OpRewritePattern<tensor::DimOp> {
----------------
MaheshRavishankar wrote:

> How about adding it here?
> 
> https://github.com/llvm/llvm-project/blob/6786928c4fe1f9daf720d3b604987de2b013e70b/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp#L132
> 
> the same pattern exists for `forop`.

This pattern is for resolving the dim of the result of the for loop. Thats is different (and should actually be done through the interface, I am not sure why it is not). 

What I am suggesting is currently the `tensor.dim` resolution through use of `ReifyRankedShapedTypeOpInterface` in [this pattern](https://github.com/llvm/llvm-project/blob/394f59c203c715d00be4643c20bbe22893397adf/mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp#L78) . We could add a new pattern that resolve the dimensions of tensor.dim operations of `iter_args` of `LoopLikeOpInterface`. You can just add that pattern to https://github.com/llvm/llvm-project/blob/394f59c203c715d00be4643c20bbe22893397adf/mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp#L127

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