[mlir][linalg] Add reduction tiling transformation

[ThomasRaoux]

Summary

Add a transformation to tile reduction ops into a parallel operation
followed by a merge operation. This is equivalent to the existing
reduction spliting transformation but using loops instead of using
higher dimensions linalg.

[mravishankar]

Wow! Awesome! I was just going to accept this, but realized I have one question on further reading. There seems to be an implicit assumption about there being a single reduction dimension through the different interface methods. Maybe its better to be more explicit, and just make the dimension (or list of dimensions) that is being tiled be an explicit argument to the interface methods. That will remove that implicit assumption through the different methods.

mlir/include/mlir/Interfaces/TilingInterface.td

163 ↗

(On Diff #470086)

Is there are reason the interface method has to return a tensor. I think if it returns the identity element that should be enough? I am not opposed to doing this.. This might actually be more general, but trying to understand what is the minimum requirement that the interface should expect ops to implement.

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

273 ↗

(On Diff #470086)

If I am reading this correctly, the setting of the identity element can be moved into the tileReductionUsingSCFForOp method (related to the comment above)

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

253

I think SCF already has a depence on Arith dialect. I think you can replace this with getValueOrCreateConstantIndexOp.

[ThomasRaoux]

In D136586#3879727, @mravishankar wrote:

Wow! Awesome! I was just going to accept this, but realized I have one question on further reading. There seems to be an implicit assumption about there being a single reduction dimension through the different interface methods. Maybe its better to be more explicit, and just make the dimension (or list of dimensions) that is being tiled be an explicit argument to the interface methods. That will remove that implicit assumption through the different methods.

Yes I can do that. I was trying to make the interface more forward looking as I think we should be able to support multiple reductions at some point but I agree it is a bit odd. I'll change the interfaces for now.

mlir/include/mlir/Interfaces/TilingInterface.td

163 ↗

(On Diff #470086)

If we return the identity element we still need to create op generating the tensor which would be dependent on the op we are tiling. (for linalg we would generate linalg.fill but not for other ops). Where would you do that?

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

273 ↗

(On Diff #470086)

Getting the reductionOp requires code depending on linalg so I'm not sure how we could do that without adding more interface functions.

[vmurali]

Flyby comment: looks like this differential is essentially doing what is being done in https://github.com/iree-org/iree/pull/10537/files#diff-e105ec843efed97f5884ad26c85561c72967f0546029bd772335634f23553c31, which means I can remove that PR and use your new interface. You seem to have gotten rid of the control for inserted index. Do you no longer need that for general purpose reduction tiling?

[ThomasRaoux]

In D136586#3880177, @vmurali wrote:

Flyby comment: looks like this differential is essentially doing what is being done in https://github.com/iree-org/iree/pull/10537/files#diff-e105ec843efed97f5884ad26c85561c72967f0546029bd772335634f23553c31, which means I can remove that PR and use your new interface. You seem to have gotten rid of the control for inserted index. Do you no longer need that for general purpose reduction tiling?

Yes it could be a potential way to replace this and generalize it by supporting dynamic shapes. The control index can be added to this if needed. Right now the code just makes a decision based on the reduction loop index but for more flexibility there should be a control indeed.

[vmurali]

In D136586#3879727, @mravishankar wrote:

Wow! Awesome! I was just going to accept this, but realized I have one question on further reading. There seems to be an implicit assumption about there being a single reduction dimension through the different interface methods. Maybe its better to be more explicit, and just make the dimension (or list of dimensions) that is being tiled be an explicit argument to the interface methods. That will remove that implicit assumption through the different methods.

I think this OP should _not_ take care of multiple reduction dimensions. There's already a separate transformation op or transformation (that Thomas wrote earlier) that combines multiple contiguous reduction dimensions - so we can use that to combine the sequence of innermost reductions into one. If the reduction dimension is not innermost, then there's no need to apply the technique as parallel tiling will already create tiles that can be reduced in parallel.

In other words, I think this differential should be accepted as is, with comments (and checks) denoting that only innermost dimension must be reduced using this transformation.

[pifon2a]

mlir/include/mlir/Interfaces/TilingInterface.td

200 ↗

(On Diff #470086)

why are mergeReductions, tileToPartialReduction, generateInitialTensorForPartialReduction a part of tiling interface? They work only for Linalg, don't they?

[Hardcode84]

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

1064

maybe return Optional<Attribute> or FailureOr<Attribute> to force users to check validity?

1079

Did you mean getLargest(semantic, false)? Also, why getLargest and not getInf?

[ThomasRaoux]

Move the reduction detection dimension in the common code and address other review comments.

[ThomasRaoux]

more review comments

[ThomasRaoux]

In D136586#3879727, @mravishankar wrote:

Wow! Awesome! I was just going to accept this, but realized I have one question on further reading. There seems to be an implicit assumption about there being a single reduction dimension through the different interface methods. Maybe its better to be more explicit, and just make the dimension (or list of dimensions) that is being tiled be an explicit argument to the interface methods. That will remove that implicit assumption through the different methods.

Thanks Mahesh. As discussed I moved the code detecting the reduction dimension to the common code but creating the tensor is still done by the interface function.

mlir/include/mlir/Interfaces/TilingInterface.td

200 ↗

(On Diff #470086)

The idea is that other ops than linalg could implement this interface in the future. For instance we did have to do reduction splitting for linalgext topK and that could be an op that would implement this interface. Also since we need to re-use helpers in the tiling code this allows us to do that without adding dependency between the tiling code and linalg.

[ThomasRaoux]

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

1064

Changed it.

1079

You are right this is a bug. Since I'm only moving this code here I would rather not add the bug fix as part of this patch. I'll send another patch to fix this independently of this change,

[herhut]

mlir/include/mlir/Interfaces/TilingInterface.td

200 ↗

(On Diff #470086)

Is this really tiling the operation? Yes, at some level it moves the computation to working on tiles but those tiles are not "normal" tiles in the sense of all other operations in this interface. Instead, these methods allow to "distribute" a reduction over multiple threads. So I see how having these methods implemented allows to apply this transformation to anything that is like a reduction but they are not useful otherwise. Maybe have a StructuredReductionOp interface?

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

315 ↗

(On Diff #470694)

Why is omitPartialTileCheck true here? Do we statically know that the dimension tiles?

mlir/test/Dialect/Linalg/transform-tile-reduction.mlir

37 ↗

(On Diff #470694)

Why is this valid? What happens if ARG0 does not have a size that is divisible by 5?

[mravishankar]

mlir/include/mlir/Interfaces/TilingInterface.td

200 ↗

(On Diff #470086)

I dont see a reason why this needs to be moved out of TilingInterface into any other interface... If an op doesnt support this kind of tiling then it just does not implement these methods. Also I think there is a final state where we can combine tileToPartialReduction and getTiledImplementation, but that isnt clear yet. So there is a separate method added for now. It is meant to handle tiling of loops that are reductions, which are fine. Haing more interfaces just adds strange dependencies, or results in having redundant interface methods across interfaces.

189 ↗

(On Diff #470694)

Could we make this an ArrayRef<unsigned>. Also I have been lead to believe that recommended C++ usage is to use int types unless you are doing bit manipulation.

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

315 ↗

(On Diff #470694)

The implementation of the loops already handles the partial tile handling. That is accounted for in the sizes. There is no need for an op implementation to account for it.

mlir/test/Dialect/Linalg/transform-tile-reduction.mlir

37 ↗

(On Diff #470694)

This is a good point. I'd expect there to be an affine.min generated by generateTiledLoopNest for this case....

[ThomasRaoux]

Fix bug in insert_slice size

[ThomasRaoux]

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

315 ↗

(On Diff #470694)

Yes, this is handled by generateTileLoopNest so the sizes are already adjusted.

mlir/test/Dialect/Linalg/transform-tile-reduction.mlir

37 ↗

(On Diff #470694)

good catch, there was a bug in the size used to insert. It is fixes now and the size is based on affine_min as expected.

[ThomasRaoux]

Address more review comments.

mlir/include/mlir/Interfaces/TilingInterface.td

189 ↗

(On Diff #470694)

Changed it to ArrayRef<int>

[ThomasRaoux]

Fix alignments

[pifon2a]

mlir/include/mlir/Interfaces/TilingInterface.td

200 ↗

(On Diff #470086)

I am worried that this change modifies TilingInterface substantially without a proper RFC and it is not clear whether these methods should be a part of the interface at all.

[mravishankar]

mlir/include/mlir/Interfaces/TilingInterface.td

200 ↗

(On Diff #470086)

This is a strict addition to the TilingInterface, and is an opt-in. So does not impact anyone who is using TilingInterface. So I dont really think it needs an RFC. Can you give more substantial reasoning as to why this should not be part of the TilingInterface? It is reusing the implementation of these methods in other operations, as well as reusing parts of the implementation of tiling using the interface. That is a strong signal to me that it indeed belongs here. I dont see a reason to block this patch. If indeed it turns out to be muddling separation of concerns we can revisit that, but to me this is already layered correctly.

[pifon2a]

mlir/include/mlir/Interfaces/TilingInterface.td

200 ↗

(On Diff #470086)

These methods are needed only for a subset of tileable operations. Adding it all here is similar to adding everything to LinalgStructuredOp interface and then spending a lot of time extracting DPS interface out of it, because LinalgOps are just a subset of DPS ops.

[mravishankar]

mlir/include/mlir/Interfaces/TilingInterface.td

200 ↗

(On Diff #470086)

I don't think they are equivalent. The DPS was deemed not a subset of LinalgStructuredOp, but rather a superset. That's why it was pulled out. As you say this is needed only for a subset of tileable operations. If an operation cannot support this way of Tiling, it just doesn't implement those interface methods. Then you cannot apply this transformation on those ops, and that is exactly how it should work.

[ThomasRaoux]

Move code out of TilingInterface

[ThomasRaoux]

mlir/include/mlir/Interfaces/TilingInterface.td

200 ↗

(On Diff #470086)

The code is moved all into linalg. This requires exposing more scf helpers.

[mravishankar]

I am not on board with exposing these utility functions. The layering of this going into TilingInterface is the right one which will avoid exposing these methods that are really implementation details. Lets hold off till we can reach a consensus.

[ThomasRaoux]

Moved new interface method to a separate interface outside of TilingInterface

[ThomasRaoux]

Based on discussions I moved the methods to a separate interface in parallel of TilingInterface. Please take another look.

[nicolasvasilache]

Bunch of suggestions for improving code and legitbility but nothing worth delaying more.

Nice stuff, thank you @ThomasRaoux !

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

652

Maybe document that your transform here overallocates for the sequential case (i.e. you could only have a tensor<5xf32> and iteratively reduce within the loop) but that this also works for the parallel case?

Do you expect this to be configurable in the future or to always use the larger allocation?

Edit: I realize the regular tiling to scf.for should do the "smaller allocation version". So now I am wondering whether we see an advantage doing it this way in the sequential case ? (i.e. is this related to @dcaballe 's vectorization comments ?)

mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h

143 ↗

(On Diff #472506)

better docs here plz

153 ↗

(On Diff #472506)

nit: merged

154 ↗

(On Diff #472506)

super nit: column

160 ↗

(On Diff #472506)

super nit: extra newline

mlir/include/mlir/Interfaces/TilingInterface.td

159 ↗

(On Diff #472506)

I like this new interface, not all ops can support partial-reduction tiling without quite more effort (e.g. gather/scatter/sort).
But a big class of ops does support this out of the box.

162 ↗

(On Diff #472506)

to tile reductions using partial reductions + merge ?

169 ↗

(On Diff #472506)

nit: methode is French :) (here and below)

179 ↗

(On Diff #472506)

nit: reductionDims (plural) since you take an ArrayRef ?

189 ↗

(On Diff #470694)

Our default type is int64_t everywhere, any concrete reason to use int (I don't think the memory usage for a few ints jusifies int)?

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

256 ↗

(On Diff #472506)

OpBuilder::InsertionGuard guard(b);

In anything that takes a builder reference, otherwise you're leaking insertion point changes across function boundaries which creates bugs that are very hard to track.

263 ↗

(On Diff #472506)

int64_t everywhere, all the time, plz (unless doing bit manipulations).
If you hit a size_t or unsigned you don't control, please static_cast<int64_t>

276 ↗

(On Diff #472506)

I don't see a new map here ?

292 ↗

(On Diff #472506)

Please use ShapedType::isDynamic.
We should be close to be able to make ShapedType::kDynamicSize private now.

310 ↗

(On Diff #472506)

insertion guard

317 ↗

(On Diff #472506)

Can we add a proper AffineMap.h/cpp helper:

AffineMap::injectResult(int64_t resultPos, AffineExpr newResult)

and just use that directly:

newMaps.back()  = newMaps.back() .injectResult(...);

318 ↗

(On Diff #472506)

Some comment here re injecting the split dimension (e0 .. ek .. en) -> (e0 .. ek, split_dim, ek+1 .. en)

327 ↗

(On Diff #472506)

Better comments plz:

Step 1. ...
Step 1. a. ...

333 ↗

(On Diff #472506)

Plz add a top-of-the-function // TODO: SubsetExtractOpInterface

355 ↗

(On Diff #472506)

insertion guard

365 ↗

(On Diff #472506)

I have a personal preference for fewer loops / less code under loops:

SmallVector<StringRef> reductionIteratorTypes(intermRank, getParallelIteratorTypeName());
reductionIteratorTypes[dimToMerge] = `();
AffineMap outputMap = AffineMap::multiDimSomethingMap().drop_result(dimToMerge); // or the equivalent APIs

367 ↗

(On Diff #472506)

unsigned -> int64_t everywhere plz

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

251

Interesting, we don't have a good way of enforcing that an interface implies another.
Can you add a top-level comment that this is assumed?

265

Can we just use an llvm::count and llvm::any_of/all_of ?
The readability improvements would greatly trump the performance gains of fusion..

int64_t numReductionDims = llvm::count/count_if(lambda);
if (numReductionDims != 1) error;

if (llvm::any_of(lambda))
  return b.notifyMatchFailure(op, "only reduction dimensions can have non zero tile sizes");

286

return b.notifyMatchFailure( for better debugging

289

Can you add a TODO at the func declaration that we should use ArrayRef<OpFoldResult> instead of ArrayRef<Value>?

294

something's fishy between this insertion point and the next .. either the next one crashes or this one does nothing.

304

You need to wrap this into something else to get a real OpFoldResult.
However, with the current broken createOrFold API you can never avoir emitting the constant op ..

312

Is the else ever legitimate ?
This needs a comment with a counter-example.

[herhut]

Thanks for refactoring. Having this in a separate interface addresses my concerns about the partial interface. Did not look at the deep details but also do not want to block this.

[mravishankar]

Nicolas already seems to have reviewed in great detail! The overall structure looks fine to me. Thanks @ThomasRaoux for going through the different iterations to reach convergence.

mlir/include/mlir/Interfaces/TilingInterface.td

159 ↗

(On Diff #472506)

Agreed! sort can actually support it I think :)

[ThomasRaoux]

Address review comments.

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

652

I'm adding more doc, in general the only case where we would want to allocate less is if the reduction size is smaller than the tile size. This could become an affin.min but I'm not sure we would ever want that. We can iterate if needed.

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

256 ↗

(On Diff #472506)

even if I don't move the insertion point?
Added to all of those right now

276 ↗

(On Diff #472506)

Yes the comment was outdated. Rewrote it.

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

251

Makes sense, added a comment.

294

oops yeah removed this code

312

I don't think there is, I removed the if

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

def SplitReductionOp : Op<Transform_Dialect, "structured.split_reduction",
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target, 
         ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
         ::mlir::transform::TransformState &state);
   }];
 }
 
+def TileReductionUsingScfOp : Op<Transform_Dialect, "structured.tile_reduction_using_scf",
+       [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+        TransformEachOpTrait, TransformOpInterface]> {
+  let description = [{
+    Indicates that the given `target` op should be transformed with the 
+    `tileReduction` transformation with the tile size provided as attribute.
+    
+    This transformation tiles the `target` along the reduction dimensions. It
+    creates a tensor initialized with the identity value. Then it creates nested
+    loops with a parallel version of `target` op inside. The parallel op
+    dimensions are less or equal to the tile size passed by user.
+    After the loop a merge operation is created to do a final reduction with the
+    partial reductions.
+
+    #### Return modes
+
+    This 3 returned handles point to:
+      - the fill op used to initialize the neutral element, 
+      - the parallel tiled op and 
+      - the result-combining op.
+
+    #### Example:
+    
+    ```
+      %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                                              affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<?x?xf32>)
+      outs(%out : tensor<?xf32>) {
+        ^bb0(%arg7: f32, %arg9: f32):
+        %1 = arith.addf %arg7, %arg9 : f32
+        linalg.yield %1 : f32
+      } -> tensor<?xf32>
+      return %red : tensor<?xf32>
+    ```
+  
+    is transformed into:
+    
+    ```
+      %0 = tensor.empty(%dim_1) : tensor<?x5xf32>
+      %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x5xf32>) -> tensor<?x5xf32>
+      %2 = scf.for %arg2 = %c0 to %dim_0 step %c5 iter_args(%arg3 = %1) -> (tensor<?x5xf32>) {
+        %extracted_slice = tensor.extract_slice %1[0, 0] [%dim, 5] [1, 1] : tensor<?x5xf32> to tensor<?x5xf32>
+        %extracted_slice_2 = tensor.extract_slice %arg0[0, %arg2] [%dim, 5] [1, 1] : tensor<?x?xf32> to tensor<?x5xf32>
+        %4 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, 
+                                              affine_map<(d0, d1) -> (d0, d1)>],
+        iterator_types = ["parallel", "parallel"]}
+        ins(%extracted_slice_2 : tensor<?x5xf32>)
+        outs(%extracted_slice : tensor<?x5xf32>) {
+        ^bb0(%in: f32, %out: f32):
+          %5 = arith.addf %in, %out : f32
+          linalg.yield %5 : f32
+        } -> tensor<?x5xf32>
+        %dim_3 = tensor.dim %1, %c0 : tensor<?x5xf32>
+        %inserted_slice = tensor.insert_slice %4 into %arg3[0, 0] [%dim_3, 5] [1, 1] : tensor<?x5xf32> into tensor<?x5xf32>
+        scf.yield %inserted_slice : tensor<?x5xf32>
+      }
+      %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                                            affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%2 : tensor<?x5xf32>)
+      outs(%arg1 : tensor<?xf32>) {
+      ^bb0(%in: f32, %out: f32):
+        %4 = arith.addf %in, %out : f32
+        linalg.yield %4 : f32
+      } -> tensor<?xf32>
+    ```
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                   DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes);
+  let results = (outs PDL_Operation:$fill_op,
+                      PDL_Operation:$split_linalg_op,
+                      PDL_Operation:$combining_linalg_op);
+
+  let assemblyFormat = "$target attr-dict";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::linalg::LinalgOp target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 def TileOp : Op<Transform_Dialect, "structured.tile",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let description = [{
     Indicates that the given `target` op should be tiled with the given sizes.
     This transform generates a loop nest with a smaller ("tiled") target
     operation in its body. Currently limited to LinalgOps.
 

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 ///
 ///  return %4 : tensor<16x32xf32>
 /// ```
 FailureOr<SplitReductionResult>
 splitReductionByScaling(PatternRewriter &b, LinalgOp op,
                         const ControlSplitReductionFn &controlSplitReductionFn,
                         bool useAlloc = false);
 
+/// Transformation information returned after reduction tiling.
+struct ReductionTilingResult {
+  /// The tiled operation generated.
+  Operation *parallelTiledOp;
+  /// The tiled operation generated.
+  Operation *mergeOp;
+  /// Initial op
+  Operation *initialOp;
+  /// The `scf.for` operations that iterate over the tiles.
+  SmallVector<scf::ForOp> loops;
+};
+
+/// Method to tile a reduction and generate a parallel op within a serial loop.
+/// Each of the partial reductions are calculated in parallel. Then after the
+/// loop all the partial reduction are mered into a final reduction.
+/// For example, for the following sequence
+///
+/// ```mlir
+/// %0 = linalg.generic %in ["parallel", "reduction"]
+///   : tensor<7x9xf32> -> tensor<7xf32>
+/// ```
+/// into:
+///
+/// ```mlir
+/// %0 = linalg.fill ... : tensor<7x4xf32>
+/// %1 = scf.for ... iter_args(%arg0 = %0)
+///   %2 = tensor.extract_slice %arg0 : tensor<7x4xf32> -> tensor<7x?xf32>
+///   %3 = tensor.extract_slice %in : tensor<7x9xf32> -> tensor<7x?xf32>
+///   %4 = linalg.generic %2, %3 ["parallel", "parallel"]
+///     : tensor<7x?xf32> -> tensor<7x?xf32>
+///   %5 = tensor.insert_slice %3, %0[0, 0] : tensor<7x4xf32>
+/// }
+/// %6 = linalg.generic %1 ["parallel", "reduction"]
+///   : tensor<7x4xf32> -> tensor<7xf32>
+/// ```
+FailureOr<ReductionTilingResult>
+tileReductionUsingScf(PatternRewriter &b, LinalgOp op,
+                      ArrayRef<OpFoldResult> tileSize);
+
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 /// of a insert_slice) operation with given offsets, and sizes to its
 /// rank-reduced version. This is only done for the cases where the size is 1
 /// and offset is 0. Strictly speaking the offset 0 is not required in general,
 /// but non-zero offsets are not handled by SPIR-V backend at this point (and
 /// potentially cannot be handled).
 Optional<SmallVector<ReassociationIndices>>
 getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes);
 
+/// Return the identity numeric value associated to the give op. Return
+/// llvm::None if there is no known neutral element.
+Optional<Attribute> getNeutralElement(Operation *op);
+
 //===----------------------------------------------------------------------===//
 // Fusion / Tiling utilities
 //===----------------------------------------------------------------------===//
 
 /// The type of loops to be generated during tiling.
 enum class LinalgTilingLoopType {
   Loops = 0,
   AffineLoops = 1,

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

 
 /// Get perfectly nested sequence of loops starting at root of loop nest
 /// (the first op being another AffineFor, and the second op - a terminator).
 /// A loop is perfectly nested iff: the first op in the loop's body is another
 /// AffineForOp, and the second op is a terminator).
 void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
                              scf::ForOp root);
 
+/// Generate an empty loop nest that represents the tiled loop nest shell.
+/// - `loopRanges` specifies the lb, ub and step of the untiled iteration space.
+/// - `tileSizeVals` is the tile sizes to use. Zero represent untiled loops.
+/// - In `offsets` and `sizes` return the multi-dimensional offset and size of
+/// the
+///   tile processed within the inner most loop.
+SmallVector<scf::ForOp> generateTileLoopNest(OpBuilder &builder, Location loc,
+                                             ArrayRef<Range> loopRanges,
+                                             ArrayRef<Value> tileSizeVals,
+                                             SmallVector<OpFoldResult> &offsets,
+                                             SmallVector<OpFoldResult> &sizes);
+
+/// For a value to be yielded (`yieldedValue`) from within a loop nest `loops`,
+/// construct the destructive update pattern that inserts the yielded
+/// value into a destination tensor provided by `initValue` at offset
+/// `tileOffsets` and size `tileSizes`. For example,
+///
+/// ```mlir
+/// scf.for %iv0 = ... {
+///   %0 = tiled_op
+/// }
+/// ```
+///
+/// is transformed to
+///
+/// ```mlir
+/// scf.for %iv0 = ... iter_args(%arg = %0) {
+///   %1 = tensor.extract_slice %arg
+///   %2 = tiled_op
+///   %3 = tensor.insert_slice %2 into %arg
+///   scf.yield %3
+/// }
+/// ```
+/// TODO: This API can be cleaned up by using `SubsetExtractOpInterface`.
+FailureOr<SmallVector<Value>>
+yieldTiledValues(RewriterBase &rewriter, ValueRange initValues,
+                 ValueRange yieldedValues,
+                 ArrayRef<SmallVector<OpFoldResult>> tileOffsetsList,
+                 ArrayRef<SmallVector<OpFoldResult>> tileSizesList,
+                 MutableArrayRef<scf::ForOp> loops);
+
+/// If the tiled operation is destination passing style, update the
+/// slice of the destination used (which refers to the untiled destination)
+/// to use the corresponding region argument of the innermost loop.
+///
+/// ```mlir
+/// %0 =
+/// scf.for %iv0 = ... iter_args(%arg = %0) {
+///   %1 = tensor.extract_slice %0
+///   %2 = tiled_op
+///   %3 = tensor.insert_slice %2 into %arg
+///   scf.yield %3
+/// }
+/// ```
+///
+/// is transformed to
+///
+/// ```mlir
+/// scf.for %iv0 = ... iter_args(%arg = %0) {
+///   %1 = tensor.extract_slice %arg
+///   %2 = tiled_op
+///   %3 = tensor.insert_slice %2 into %arg
+///   scf.yield %3
+/// }
+/// ```
+void updateDestinationOperandsForTiledOp(OpBuilder &builder,
+                                         ValueRange tiledOpDestinationValues,
+                                         ValueRange bbArgsList);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_UTILS_UTILS_H_

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

transform::SplitReductionOp::applyToOne(linalg::LinalgOp target,
   results.push_back(splitResult->initOrAlloc);
   results.push_back(splitResult->fillOp);
   results.push_back(splitResult->splitLinalgOp);
   results.push_back(splitResult->resultCombiningLinalgOp);
   return DiagnosedSilenceableFailure(success());
 }
 
 //===----------------------------------------------------------------------===//
+// SplitReductionOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::TileReductionUsingScfOp::applyToOne(
+    linalg::LinalgOp target, SmallVectorImpl<Operation *> &results,
+    transform::TransformState &state) {
+  SimpleRewriter rewriter(getContext());
+  rewriter.setInsertionPoint(target);
+  SmallVector<int64_t> tileSizes = extractFromI64ArrayAttr(getTileSizes());
+  SmallVector<OpFoldResult> sizes;
+  for (int64_t size : tileSizes) {
+    sizes.push_back(rewriter.getIndexAttr(size));
+  }
+
+  FailureOr<linalg::ReductionTilingResult> result =
+      linalg::tileReductionUsingScf(rewriter, target, sizes);
+
+  if (failed(result))
+    return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
+  results.push_back(result->initialOp);
+  results.push_back(result->parallelTiledOp);
+  results.push_back(result->mergeOp);
+  return DiagnosedSilenceableFailure(success());
+}
+
+//===----------------------------------------------------------------------===//
 // TileOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::TileOp::apply(TransformResults &transformResults,
                          TransformState &state) {
   SmallVector<int64_t> tileSizes = extractFromI64ArrayAttr(getStaticSizes());
 

mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp

 // between a parallel and a reduction dimension.
 //
 //===----------------------------------------------------------------------===//
 
 #include <utility>
 
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
-/// Return the identity numeric value associated to the give op.
-static Attribute getNeutralElement(Operation *op) {
-  // Builder only used as helper for attribute creation.
-  OpBuilder b(op->getContext());
-  Type resultType = op->getResult(0).getType();
-  if (auto floatType = resultType.dyn_cast<FloatType>()) {
-    const llvm::fltSemantics &semantic = floatType.getFloatSemantics();
-    if (isa<arith::AddFOp>(op))
-      return b.getFloatAttr(resultType, llvm::APFloat::getZero(semantic));
-    if (isa<arith::MulFOp>(op))
-      return b.getFloatAttr(resultType, llvm::APFloat(semantic, 1));
-    if (isa<arith::MaxFOp>(op))
-      return b.getFloatAttr(resultType,
-                            llvm::APFloat::getLargest(semantic, true));
-    if (isa<arith::MinFOp>(op))
-      return b.getFloatAttr(resultType,
-                            llvm::APFloat::getLargest(semantic, true));
-    return Attribute();
-  }
-  if (isa<arith::AddIOp, arith::OrIOp, arith::XOrIOp>(op))
-    return b.getIntegerAttr(resultType, 0);
-  if (isa<arith::AndIOp>(op))
-    return b.getIntegerAttr(resultType, -1);
-  if (isa<arith::MaxSIOp>(op))
-    return b.getIntegerAttr(resultType, std::numeric_limits<int64_t>::min());
-  if (isa<arith::MinSIOp>(op))
-    return b.getIntegerAttr(resultType, std::numeric_limits<int64_t>::max());
-  if (isa<arith::MulIOp>(op))
-    return b.getIntegerAttr(resultType, 1);
-  return Attribute();
-}
-
 FailureOr<SplitReductionResult> mlir::linalg::splitReduction(
     PatternRewriter &b, LinalgOp op,
     const ControlSplitReductionFn &controlSplitReductionFn, bool useAlloc) {
   OpBuilder::InsertionGuard guard(b);
   b.setInsertionPoint(op);
 
   SplitReductionOptions control = controlSplitReductionFn(op);
   int64_t ratio = control.ratio;
return b.notifyMatchFailure(
         op, "Reduction dimension not divisible by split ratio");
 
   SmallVector<Operation *, 4> combinerOps;
   if (!matchReduction(op.getRegionOutputArgs(), 0, combinerOps) ||
       combinerOps.size() != 1)
     return b.notifyMatchFailure(op, "Cannot match the reduction pattern");
 
   Operation *reductionOp = combinerOps[0];
-  Attribute identity = getNeutralElement(reductionOp);
-  if (!identity)
+  Optional<Attribute> identity = getNeutralElement(reductionOp);
+  if (!identity.has_value())
     return b.notifyMatchFailure(op, "Unknown identity value for the reduction");
 
   Location loc = op->getLoc();
   SmallVector<Value> newInputs;
   SmallVector<AffineMap> newMaps;
   // Calculate the new shapes and indexing maps of the input operands.
   for (OpOperand *operand : op.getInputOperands()) {
     AffineMap map = op.getMatchingIndexingMap(operand);
emptyOrAllocTensor = b.create<bufferization::AllocTensorOp>(
         loc,
         RankedTensorType::get(newOutputShape,
                               op.getRegionOutputArgs()[0].getType()),
         ValueRange{});
   } else {
     emptyOrAllocTensor = b.create<tensor::EmptyOp>(
         loc, newOutputShape, op.getRegionOutputArgs()[0].getType());
   }
-  Value constantOp = b.create<arith::ConstantOp>(loc, identity);
+  Value constantOp = b.create<arith::ConstantOp>(loc, *identity);
   Value identityTensor =
       b.create<linalg::FillOp>(op->getLoc(), constantOp, emptyOrAllocTensor)
           .getResult(0);
 
   newMaps.push_back(AffineMap::get(oldOutputMap.getNumDims() + 1, 0, outputExpr,
                                    op.getContext()));
   SmallVector<StringRef> newIteratorTypes;
   for (auto &it : llvm::enumerate(op.getIteratorTypesArray())) {
if (reductionDimSize == ShapedType::kDynamicSize ||
       insertSplitDimension >= loopRanges.size())
     return b.notifyMatchFailure(
         op, "first reduction dimension not divisible by split factor");
 
   SmallVector<Operation *> combinerOps;
   if (!matchReduction(op.getRegionOutputArgs(), 0, combinerOps))
     return b.notifyMatchFailure(op, "cannot match a reduction pattern");
 
-  SmallVector<Attribute> neutralElements = llvm::to_vector<4>(
-      llvm::map_range(combinerOps, [&](Operation *reductionOp) {
-        return getNeutralElement(reductionOp);
-      }));
+  SmallVector<Attribute> neutralElements;
+  for (Operation *reductionOp : combinerOps) {
+    Optional<Attribute> neutralElement = getNeutralElement(reductionOp);
+    if (!neutralElement.has_value())
+      return b.notifyMatchFailure(op, "cannot find neutral element.");
+    neutralElements.push_back(*neutralElement);
+  }
   if (!llvm::all_of(neutralElements, [](Attribute attr) { return attr; }))
     return b.notifyMatchFailure(op, "unknown reduction neutral");
 
   // TODO: relax this when multi-reduction support is available.
   if (op.getNumOutputs() != static_cast<int64_t>(neutralElements.size()))
     return b.notifyMatchFailure(op, "expect one reduction per output");
 
   // Rewrite part.
FailureOr<SplitReductionResult> mlir::linalg::splitReductionByScaling(
   // TODO: extend when multi-reduction support is available.
   assert(fillOps.size() == results.size() && results.size() == 1);
   b.replaceOp(op, results.front()->getResults());
   return SplitReductionResult{emptyOrAllocTensorOps.front(), fillOps.front(),
                               cast<LinalgOp>(genericOp.getOperation()),
                               results.front()};
 }
 
+static FailureOr<Operation *> generateInitialTensorForPartialReduction(
+    Operation *op, OpBuilder &b, Location loc, ArrayRef<OpFoldResult> sizes,
+    ArrayRef<int> reductionDims) {
+  auto linalgOp = cast<LinalgOp>(op);
+  assert(reductionDims.size() == 1 &&
+         "only support single reduction right now.");
+  if (linalgOp.hasBufferSemantics())
+    return op->emitOpError("expected operation to have tensor semantics");
+  // Insert the new parallel dimension based on the index of the reduction
+  // loop. This could be controlled by user for more flexibility.
+  unsigned insertSplitDimension = unsigned(reductionDims[0]);
+
+  SmallVector<Operation *, 4> combinerOps;
+  if (!matchReduction(linalgOp.getRegionOutputArgs(), 0, combinerOps) ||
+      combinerOps.size() != 1)
+    return op->emitOpError("Failed to anaysis the reduction operation.");
+
+  Operation *reductionOp = combinerOps[0];
+  Optional<Attribute> identity = getNeutralElement(reductionOp);
+  if (!identity.has_value())
+    return op->emitOpError(
+        "Failed to get an identity value for the reduction operation.");
+
+  // Calculate the new output map and shape, we insert the new dimension in
+  // the intermediate tensor based on the index returned by
+  // `controlSplitReductionFn`.
+  SmallVector<int64_t> newOutputShape;
+  ArrayRef<int64_t> oldShape = linalgOp.getShape(linalgOp.getOutputOperand(0));
+  SmallVector<Value> dynamicDims;
+  for (unsigned idx : llvm::seq<unsigned>(0, oldShape.size() + 1)) {
+    if (idx == insertSplitDimension) {
+      dispatchIndexOpFoldResults(sizes[idx], dynamicDims, newOutputShape,
+                                 ShapedType::kDynamicStrideOrOffset);
+      continue;
+    }
+    unsigned oldIdx = idx < insertSplitDimension ? idx : idx - 1;
+    int64_t dim = oldShape[oldIdx];
+    newOutputShape.push_back(dim);
+    if (dim == ShapedType::kDynamicSize)
+      dynamicDims.push_back(b.createOrFold<tensor::DimOp>(
+          loc, linalgOp.getOutputOperand(0)->get(), oldIdx));
+  }
+  Value emptyTensor = b.create<tensor::EmptyOp>(
+      loc, newOutputShape, linalgOp.getRegionOutputArgs()[0].getType(),
+      dynamicDims);
+  Value constantOp = b.create<arith::ConstantOp>(loc, *identity);
+  auto identityTensor = b.create<linalg::FillOp>(loc, constantOp, emptyTensor);
+  return identityTensor.getOperation();
+}
+
+static Operation *tileToPartialReduction(Operation *op, OpBuilder &b,
+                                         Location loc, ValueRange init,
+                                         ArrayRef<OpFoldResult> offsets,
+                                         ArrayRef<OpFoldResult> sizes,
+                                         ArrayRef<int> reductionDims) {
+  auto linalgOp = cast<LinalgOp>(op);
+  assert(reductionDims.size() == 1 &&
+         "only support single reduction right now.");
+  unsigned insertSplitDimension = unsigned(reductionDims[0]);
+
+  AffineMap oldOutputMap =
+      linalgOp.getMatchingIndexingMap(linalgOp.getOutputOperand(0));
+  SmallVector<AffineExpr> outputExpr;
+  for (auto &[idx, expr] : llvm::enumerate(oldOutputMap.getResults())) {
+    if (idx == insertSplitDimension) {
+      outputExpr.push_back(b.getAffineDimExpr(reductionDims[0]));
+    }
+    outputExpr.push_back(expr);
+  }
+  if (insertSplitDimension == oldOutputMap.getNumResults())
+    outputExpr.push_back(b.getAffineDimExpr(reductionDims[0]));
+
+  SmallVector<Value> valuesToTile = linalgOp.getInputOperands();
+  SmallVector<Value, 4> tiledOperands =
+      makeTiledShapes(b, loc, op, valuesToTile, offsets, sizes, {}, true);
+
+  SmallVector<OpFoldResult> strides(offsets.size(), b.getIndexAttr(1));
+  SmallVector<OpFoldResult> outOffsets(offsets.size(), b.getIndexAttr(0));
+  Value out = b.create<tensor::ExtractSliceOp>(loc, init[0], outOffsets, sizes,
+                                               strides);
+
+  // Create the new op matching the original op where the reduction dimension
+  // is now parallel.
+  SmallVector<StringRef> newIteratorTypes = linalgOp.getIteratorTypesArray();
+  newIteratorTypes[reductionDims[0]] = getParallelIteratorTypeName();
+  SmallVector<AffineMap> newMaps = linalgOp.getIndexingMapsArray();
+  newMaps.back() = AffineMap::get(newMaps.back().getNumDims(), 0, outputExpr,
+                                  linalgOp.getContext());
+  auto genericOp =
+      b.create<GenericOp>(loc, TypeRange({out.getType()}), tiledOperands,
+                          ValueRange({out}), newMaps, newIteratorTypes);
+  BlockAndValueMapping mapping;
+  op->getRegion(0).cloneInto(&genericOp.getRegion(),
+                             genericOp.getRegion().begin(), mapping);
+  return genericOp.getOperation();
+}
+
+static Operation *mergeReductions(Operation *op, OpBuilder &b, Location loc,
+                                  ValueRange partialReduce,
+                                  ArrayRef<int> reductionDims) {
+  auto linalgOp = cast<LinalgOp>(op);
+  assert(reductionDims.size() == 1 &&
+         "only support single reduction right now.");
+  unsigned dimToMerge = reductionDims[0];
+
+  // Then create a new reduction that only reduce the newly added dimension
+  // from the previous op.
+  unsigned intermRank = partialReduce[0].getType().cast<ShapedType>().getRank();
+  AffineMap inputMap = b.getMultiDimIdentityMap(intermRank);
+  SmallVector<StringRef> reductionIteratorTypes;
+  SmallVector<AffineExpr> exprs;
+  for (unsigned i : llvm::seq<unsigned>(0, intermRank)) {
+    if (dimToMerge == i) {
+      reductionIteratorTypes.push_back(getReductionIteratorTypeName());
+    } else {
+      exprs.push_back(b.getAffineDimExpr(i));
+      reductionIteratorTypes.push_back(getParallelIteratorTypeName());
+    }
+  }
+  AffineMap outputMap = AffineMap::get(intermRank, 0, exprs, op->getContext());
+  SmallVector<AffineMap> reductionMaps = {inputMap, outputMap};
+
+  SmallVector<Operation *, 4> combinerOps;
+  matchReduction(linalgOp.getRegionOutputArgs(), 0, combinerOps);
+  Operation *reductionOp = combinerOps[0];
+
+  auto reduction = b.create<GenericOp>(
+      loc, op->getResultTypes(), ValueRange({partialReduce[0]}),
+      SmallVector<Value>{linalgOp.getOutputOperands()}, reductionMaps,
+      reductionIteratorTypes,
+      [reductionOp](OpBuilder &b, Location loc, ValueRange inputs) {
+        Operation *clonedReductionOp = b.clone(*reductionOp);
+        clonedReductionOp->setOperand(0, inputs[0]);
+        clonedReductionOp->setOperand(1, inputs[1]);
+        b.create<linalg::YieldOp>(loc, clonedReductionOp->getResult(0));
+      });
+  return reduction.getOperation();
+}
+
+FailureOr<linalg::ReductionTilingResult>
+mlir::linalg::tileReductionUsingScf(PatternRewriter &b, LinalgOp op,
+                                    ArrayRef<OpFoldResult> tileSize) {
+  Location loc = op.getLoc();
+  auto tilingInterfaceOp = cast<TilingInterface>(op.getOperation());
+  SmallVector<Range> iterationDomain = tilingInterfaceOp.getIterationDomain(b);
+  SmallVector<Value> tileSizeVector =
+      getValueOrCreateConstantIndexOp(b, loc, tileSize);
+  if (tileSizeVector.size() < iterationDomain.size()) {
+    auto zero = b.create<arith::ConstantIndexOp>(loc, 0);
+    tileSizeVector.append(iterationDomain.size() - tileSizeVector.size(), zero);
+  }
+  if (op->getNumResults() != 1)
+    return b.notifyMatchFailure(
+        op, "don't support ops with multiple results for now");
+  Optional<int> reductionDim;
+  for (auto &[idx, iteratorType] :
+       llvm::enumerate(tilingInterfaceOp.getLoopIteratorTypes())) {
+    if (iteratorType == utils::IteratorType::reduction) {
+      // Only support cases with one reduction dimension for now.
+      if (reductionDim.has_value())
+        return failure();
+      reductionDim = idx;
+      continue;
+    }
+    if (!isConstantIntValue(tileSize[idx], 0))
+      return b.notifyMatchFailure(
+          op, "only reduction dimensions can have non zero tile sizes");
+  }
+  if (!reductionDim.has_value())
+    return b.notifyMatchFailure(op, "doesn't have a reduction dimension");
+  // 1. create the inital tensor value.
+  FailureOr<Operation *> identityTensor =
+      generateInitialTensorForPartialReduction(op, b, loc, tileSize,
+                                               *reductionDim);
+
+  if (failed(identityTensor))
+    return failure();
+  // 2. Create the nested loops.
+  SmallVector<OpFoldResult> offsets, sizes;
+  SmallVector<scf::ForOp> loops = generateTileLoopNest(
+      b, loc, iterationDomain, tileSizeVector, offsets, sizes);
+
+  // 3. Generate the tiled implementation within the inner most loop.
+  if (!loops.empty())
+    b.setInsertionPoint(loops.back().getBody()->getTerminator());
+
+  b.setInsertionPoint(loops.back().getBody()->getTerminator());
+  Operation *parallelOp =
+      tileToPartialReduction(op, b, loc, identityTensor.value()->getResults(),
+                             offsets, sizes, *reductionDim);
+
+  SmallVector<OpFoldResult> resultSizesList;
+  for (size_t i = 0; i < offsets.size(); i++)
+    resultSizesList.push_back(
+        b.createOrFold<tensor::DimOp>(loc, parallelOp->getResult(0), i));
+  SmallVector<OpFoldResult> outOffsets(offsets.size(), b.getIndexAttr(0));
+  FailureOr<SmallVector<Value>> replacementOr = yieldTiledValues(
+      b, identityTensor.value()->getResults(), parallelOp->getResults(),
+      outOffsets, resultSizesList, loops);
+  if (failed(replacementOr))
+    return b.notifyMatchFailure(op, "failed to yield replacement");
+
+  if (auto dstOp = dyn_cast<DestinationStyleOpInterface>(parallelOp)) {
+    auto innerMostLoop = loops.back();
+    SmallVector<Value> destinationTensors = dstOp.getOutputOperands();
+    assert(destinationTensors.size() ==
+               innerMostLoop.getRegionIterArgs().size() &&
+           "unexpected number of outputs");
+    updateDestinationOperandsForTiledOp(b, destinationTensors,
+                                        innerMostLoop.getRegionIterArgs());
+  }
+
+  // 4. Apply the merge reduction to combine all the partial values.
+  b.setInsertionPointAfter(*loops.begin());
+  Operation *mergeOp =
+      mergeReductions(op, b, loc, replacementOr.value(), *reductionDim);
+  b.replaceOp(op, mergeOp->getResults());
+
+  ReductionTilingResult results;
+  results.initialOp = identityTensor.value();
+  results.loops = std::move(loops);
+  results.parallelTiledOp = parallelOp;
+  results.mergeOp = mergeOp;
+  return results;
+}
+
 namespace {
 
 struct LinalgSplitReduction : public OpInterfaceRewritePattern<LinalgOp> {
   /// Construct a generic pattern applied to all LinalgOp that verify `filter`.
   LinalgSplitReduction(MLIRContext *context,
                        ControlSplitReductionFn controlSplitReductionFn,
                        bool useAlloc = false, PatternBenefit benefit = 1)
       : OpInterfaceRewritePattern<LinalgOp>(context, benefit),

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

computeAllSliceParameters(OpBuilder &builder, Location loc, LinalgOp linalgOp,
 
   // Construct (potentially temporary) mins and maxes on which to apply maps
   // that define tile subshapes.
   SmallVector<OpFoldResult> lbs =
       computeTileOffsets(builder, loc, ivs, tileSizes);
   SmallVector<OpFoldResult> subShapeSizes =
       computeTileSizes(builder, loc, tileSizes, sizeBounds);
 
-  assert(static_cast<int64_t>(valuesToTile.size()) ==
+  assert(static_cast<int64_t>(valuesToTile.size()) <=
              linalgOp->getNumOperands() &&
-         "expected one value to tile for every operand");
+         "more value to tile than operands.");
   SmallVector<Optional<SliceParameters>> allSliceParams;
   allSliceParams.reserve(valuesToTile.size());
-  for (OpOperand &opOperand : linalgOp->getOpOperands()) {
-    Value shapedOp = valuesToTile[opOperand.getOperandNumber()];
+  for (auto [opOperand, val] :
+       llvm::zip(linalgOp->getOpOperands(), valuesToTile)) {
+    Value shapedOp = val;
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
     AffineMap map = linalgOp.getMatchingIndexingMap(&opOperand);
     // Use `opOperand` as is if it is not tiled and not an output tensor. Having
     // an extract/insert slice pair for all output tensors simplifies follow up
     // transformations such as padding and bufferization since the
     // extract/insert slice pairs make the accessed iteration argument
     // subdomains explicit.
 
getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes) {
   // When the reassociations are not empty, then fold the remaining
   // unit-dimensions into the last dimension.  If the reassociations so far is
   // empty, then leave it emtpy. This will fold everything to a rank-0 tensor.
   if (!curr.empty() && !reassociation.empty())
     reassociation.back().append(curr.begin(), curr.end());
   return reassociation;
 }
 
+/// Return the identity numeric value associated to the give op.
+Optional<Attribute> getNeutralElement(Operation *op) {
+  // Builder only used as helper for attribute creation.
+  OpBuilder b(op->getContext());
+  Type resultType = op->getResult(0).getType();
+  if (auto floatType = resultType.dyn_cast<FloatType>()) {
+    const llvm::fltSemantics &semantic = floatType.getFloatSemantics();
+    if (isa<arith::AddFOp>(op))
+      return b.getFloatAttr(resultType, llvm::APFloat::getZero(semantic));
+    if (isa<arith::MulFOp>(op))
+      return b.getFloatAttr(resultType, llvm::APFloat(semantic, 1));
+    if (isa<arith::MaxFOp>(op))
+      return b.getFloatAttr(resultType,
+                            llvm::APFloat::getLargest(semantic, true));
+    if (isa<arith::MinFOp>(op))
+      return b.getFloatAttr(resultType,
+                            llvm::APFloat::getLargest(semantic, true));
+    return Attribute();
+  }
+  if (isa<arith::AddIOp, arith::OrIOp, arith::XOrIOp>(op))
+    return b.getIntegerAttr(resultType, 0);
+  if (isa<arith::AndIOp>(op))
+    return b.getIntegerAttr(resultType, -1);
+  if (isa<arith::MaxSIOp>(op))
+    return b.getIntegerAttr(resultType, std::numeric_limits<int64_t>::min());
+  if (isa<arith::MinSIOp>(op))
+    return b.getIntegerAttr(resultType, std::numeric_limits<int64_t>::max());
+  if (isa<arith::MulIOp>(op))
+    return b.getIntegerAttr(resultType, 1);
+  return llvm::None;
+}
+
 } // namespace linalg
 } // namespace mlir

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

static bool isPermutation(ArrayRef<int64_t> interchange) {
   }
   return seenVals.size() == interchange.size();
 }
 
 //===----------------------------------------------------------------------===//
 // tileUsingSCFForOp implementation.
 //===----------------------------------------------------------------------===//
 
-// Check if `stride` evenly divides the trip count `size - offset`.
-static bool tileDividesIterationDomain(Range loopRange) {
-  Optional<int64_t> offsetAsInt = getConstantIntValue(loopRange.offset);
-  if (!offsetAsInt)
-    return false;
-  Optional<int64_t> sizeAsInt = getConstantIntValue(loopRange.size);
-  if (!sizeAsInt)
-    return false;
-  Optional<int64_t> strideAsInt = getConstantIntValue(loopRange.stride);
-  if (!strideAsInt)
-    return false;
-  return ((sizeAsInt.value() - offsetAsInt.value()) % strideAsInt.value() == 0);
-}
-
-/// Generate an empty loop nest that represents the tiled loop nest shell.
-/// - `loopRanges` specifies the lb, ub and step of the untiled iteration space.
-/// - `tileSizeVals` is the tile sizes to use. Zero represent untiled loops.
-/// - In `offsets` and `sizes` return the multi-dimensional offset and size of
-/// the
-///   tile processed within the inner most loop.
-static SmallVector<scf::ForOp>
-generateTileLoopNest(OpBuilder &builder, Location loc,
-                     ArrayRef<Range> loopRanges, ArrayRef<Value> tileSizeVals,
-                     SmallVector<OpFoldResult> &offsets,
-                     SmallVector<OpFoldResult> &sizes) {
-  assert(!loopRanges.empty() && "expected at least one loop range");
-  assert(loopRanges.size() == tileSizeVals.size() &&
-         "expected as many tile sizes as loop ranges");
-  OpBuilder::InsertionGuard guard(builder);
-  SmallVector<scf::ForOp> loops;
-  offsets.resize(loopRanges.size());
-  sizes.resize(loopRanges.size());
-
-  // The tile size to use (to avoid out of bounds access) is  minimum of
-  // `tileSize` and `ub - iv`, where `iv` is the induction variable
-  // of the tiled loop.
-  AffineExpr s0, s1, d0;
-  bindDims(builder.getContext(), d0);
-  bindSymbols(builder.getContext(), s0, s1);
-  AffineMap minMap = AffineMap::get(1, 2, {s0, s1 - d0}, builder.getContext());
-
-  for (auto loopRange : llvm::enumerate(loopRanges)) {
-    Value offset =
-        getValueOrCreateConstantIndexOp(builder, loc, loopRange.value().offset);
-    Value size =
-        getValueOrCreateConstantIndexOp(builder, loc, loopRange.value().size);
-    // No loops if tile size is zero. Set offset and size to the loop
-    // offset and size.
-    if (matchPattern(tileSizeVals[loopRange.index()], m_Zero())) {
-      offsets[loopRange.index()] = offset;
-      sizes[loopRange.index()] = size;
-      continue;
-    }
-
-    auto loop = builder.create<scf::ForOp>(
-        loc, offset, size, tileSizeVals[loopRange.index()], ValueRange{},
-        [&](OpBuilder &bodyBuilder, Location bodyLoc, Value iv,
-            ValueRange /*iterArgs*/) {
-          bool canAvoidMap = tileDividesIterationDomain(
-              Range{loopRange.value().offset, loopRange.value().size,
-                    tileSizeVals[loopRange.index()]});
-          Value boundedTileSize =
-              (canAvoidMap)
-                  ? tileSizeVals[loopRange.index()]
-                  : builder.create<AffineMinOp>(
-                        bodyLoc, minMap,
-                        ValueRange{iv, tileSizeVals[loopRange.index()], size});
-          sizes[loopRange.index()] = boundedTileSize;
-          builder.create<scf::YieldOp>(loc);
-        });
-    offsets[loopRange.index()] = loop.getInductionVar();
-    loops.push_back(loop);
-    builder.setInsertionPoint(loop.getBody()->getTerminator());
-  }
-  return loops;
-}
-
-/// For a value to be yielded (`yieldedValue`) from within a loop nest `loops`,
-/// construct the destructive update pattern that inserts the yielded
-/// value into a destination tensor provided by `initValue` at offset
-/// `tileOffsets` and size `tileSizes`. For example,
-///
-/// ```mlir
-/// scf.for %iv0 = ... {
-///   %0 = tiled_op
-/// }
-/// ```
-///
-/// is transformed to
-///
-/// ```mlir
-/// scf.for %iv0 = ... iter_args(%arg = %0) {
-///   %1 = tensor.extract_slice %arg
-///   %2 = tiled_op
-///   %3 = tensor.insert_slice %2 into %arg
-///   scf.yield %3
-/// }
-/// ```
-/// TODO: This API can be cleaned up by using `SubsetExtractOpInterface`.
-static FailureOr<SmallVector<Value>>
-yieldTiledValues(RewriterBase &rewriter, ValueRange initValues,
-                 ValueRange yieldedValues,
-                 ArrayRef<SmallVector<OpFoldResult>> tileOffsetsList,
-                 ArrayRef<SmallVector<OpFoldResult>> tileSizesList,
-                 MutableArrayRef<scf::ForOp> loops) {
-  NewYieldValueFn yieldValueFn =
-      [&](OpBuilder &b, Location loc,
-          ArrayRef<BlockArgument> newBBArgs) -> SmallVector<Value> {
-    SmallVector<Value> inserts;
-    for (const auto &yieldedValue : llvm::enumerate(yieldedValues)) {
-      ArrayRef<OpFoldResult> tileOffsets =
-          tileOffsetsList[yieldedValue.index()];
-      ArrayRef<OpFoldResult> tileSizes = tileSizesList[yieldedValue.index()];
-      SmallVector<OpFoldResult> tileStrides(tileOffsets.size(),
-                                            b.getIndexAttr(1));
-      Value insert = b.create<tensor::InsertSliceOp>(
-          loc, yieldedValue.value(), newBBArgs[yieldedValue.index()],
-          tileOffsets, tileSizes, tileStrides);
-      inserts.push_back(insert);
-    }
-    return inserts;
-  };
-
-  SmallVector<scf::ForOp> newLoops =
-      replaceLoopNestWithNewYields(rewriter, loops, initValues, yieldValueFn,
-                                   /*replaceIterOperandsUsesInLoop =*/false);
-  for (const auto &loop : llvm::enumerate(loops)) {
-    rewriter.eraseOp(loop.value());
-    loops[loop.index()] = newLoops[loop.index()];
-  }
-  return llvm::to_vector(llvm::map_range(
-      loops.front().getResults().take_back(yieldedValues.size()),
-      [](OpResult r) -> Value { return r; }));
-}
-
-/// If the tiled operation is destination passing style, update the
-/// slice of the destination used (which refers to the untiled destination)
-/// to use the corresponding region argument of the innermost loop.
-///
-/// ```mlir
-/// %0 =
-/// scf.for %iv0 = ... iter_args(%arg = %0) {
-///   %1 = tensor.extract_slice %0
-///   %2 = tiled_op
-///   %3 = tensor.insert_slice %2 into %arg
-///   scf.yield %3
-/// }
-/// ```
-///
-/// is transformed to
-///
-/// ```mlir
-/// scf.for %iv0 = ... iter_args(%arg = %0) {
-///   %1 = tensor.extract_slice %arg
-///   %2 = tiled_op
-///   %3 = tensor.insert_slice %2 into %arg
-///   scf.yield %3
-/// }
-/// ```
-static void
-updateDestinationOperandsForTiledOp(OpBuilder &builder,
-                                    ValueRange tiledOpDestinationValues,
-                                    ValueRange bbArgsList) {
-  for (const auto &destValue : llvm::enumerate(tiledOpDestinationValues)) {
-    auto sliceOp = destValue.value().getDefiningOp<tensor::ExtractSliceOp>();
-    if (!sliceOp)
-      continue;
-    sliceOp.setOperand(0, bbArgsList[destValue.index()]);
-  }
-}
-
 /// Implementation of tiling transformation of `op` that implements the
 /// `TilingInterface` using `scf.for` to iterate over the tiles.
 FailureOr<scf::SCFTilingResult>
 mlir::scf::tileUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
                              const scf::SCFTilingOptions &options) {
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPointAfter(op);
 
mlir::scf::tileUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
   return tilingResult;
 }
 
 //===----------------------------------------------------------------------===//
 // tileConsumerAndFuseProducerGreedilyUsingSCFForOp implementation.
 //===----------------------------------------------------------------------===//
 
 /// Return the untiled producer whose slice is used in a tiled consumer. The
 /// method traverses the tile loop nest (`loops`) if needed, and returns the
 /// `iter_args` of the outer most that is encountered. Traversing the iter_args
 /// indicates that this is a destination operand of the consumer. If there was
 /// no loop traversal needed, the second value of the returned tuple is empty.
 static std::tuple<OpResult, Optional<OpOperand *>>
 getUntiledProducerFromSliceSource(OpOperand *source,
                                   ArrayRef<scf::ForOp> loops) {
   Optional<OpOperand *> destinationIterArg;
   auto loopIt = loops.rbegin();
   while (auto iterArg = source->get().dyn_cast<BlockArgument>()) {
     scf::ForOp loop = *loopIt;
     if (iterArg.getOwner()->getParentOp() != loop)
       break;
     source = &loop.getOpOperandForRegionIterArg(iterArg);
     loopIt++;
   }
   if (loopIt == loops.rend())
     destinationIterArg = source;
   return {source->get().dyn_cast<OpResult>(), destinationIterArg};
 }
 
 /// Implementation of tile consumer and fuse producer greedily.
 FailureOr<scf::SCFTileAndFuseResult>
 mlir::scf::tileConsumerAndFuseProducerGreedilyUsingSCFForOp(
     RewriterBase &rewriter, TilingInterface consumer,
     const scf::SCFTileAndFuseOptions &options) {
   // This transformation is only valid for ops that return values (i.e. not
   // valid to use with operations that have memref operands).
   if (!consumer->getNumResults()) {
     return rewriter.notifyMatchFailure(
         consumer, "invalid pattern for op with no results");
   }
 
   // 1. First tile the consumer.
   scf::SCFTileAndFuseResult tileAndFuseResult;
   llvm::SmallDenseMap<Value, int64_t> yieldedValueToResultNumber;
   {
     FailureOr<scf::SCFTilingResult> tilingResult =
         tileUsingSCFForOp(rewriter, consumer, options.tilingOptions);
     if (failed(tilingResult))
       return rewriter.notifyMatchFailure(consumer, "failed to tile consumer");
     tileAndFuseResult.tiledAndFusedOps.insert(tilingResult->tiledOp);
     tileAndFuseResult.loops = std::move(tilingResult->loops);
     for (const auto &result : llvm::enumerate(
              llvm::zip(consumer->getResults(), tilingResult->replacements))) {
       tileAndFuseResult.replacements[std::get<0>(result.value())] =
           std::get<1>(result.value());
       yieldedValueToResultNumber[tilingResult->tiledOp->getResult(
           result.index())] = result.index();
     }
   }
 
   // If there are no loops generated, fusion is immaterial.
   if (tileAndFuseResult.loops.empty())
     return tileAndFuseResult;
 
   // 2. Typically, the operands of the tiled operation are slices of the
   //    operands of the untiled operation. These are expressed in IR using
   //    `tensor.extract_slice` operations with source being the operands of the
   //    untiled operation. Create a worklist of these `tensor.extract_slice`
   //    operations. If the producers of the source of the `tensor.extract_slice`
   //    can be tiled such that the tiled value is generated in-place, that
   //    effectively tiles + fuses the operations.
   auto addCandidateSlices = [](Operation *fusedOp,
                                std::deque<tensor::ExtractSliceOp> &candidates) {
     for (Value operand : fusedOp->getOperands())
       if (auto sliceOp = operand.getDefiningOp<tensor::ExtractSliceOp>())
         candidates.push_back(sliceOp);
   };
 

mlir/lib/Dialect/SCF/Utils/Utils.cpp

 //===- Utils.cpp ---- Misc utilities for loop transformation ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements miscellaneous loop transformation routines.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
TileLoops mlir::extractFixedOuterLoops(scf::ForOp rootForOp,
 
   // TODO: for now we just ignore the result of band isolation.
   // In the future, mapping decisions may be impacted by the ability to
   // isolate perfectly nested bands.
   (void)tryIsolateBands(tileLoops);
 
   return tileLoops;
 }
+
+// Check if `stride` evenly divides the trip count `size - offset`.
+static bool tileDividesIterationDomain(Range loopRange) {
+  Optional<int64_t> offsetAsInt = getConstantIntValue(loopRange.offset);
+  if (!offsetAsInt)
+    return false;
+  Optional<int64_t> sizeAsInt = getConstantIntValue(loopRange.size);
+  if (!sizeAsInt)
+    return false;
+  Optional<int64_t> strideAsInt = getConstantIntValue(loopRange.stride);
+  if (!strideAsInt)
+    return false;
+  return ((sizeAsInt.value() - offsetAsInt.value()) % strideAsInt.value() == 0);
+}
+
+SmallVector<scf::ForOp> mlir::generateTileLoopNest(
+    OpBuilder &builder, Location loc, ArrayRef<Range> loopRanges,
+    ArrayRef<Value> tileSizeVals, SmallVector<OpFoldResult> &offsets,
+    SmallVector<OpFoldResult> &sizes) {
+  assert(!loopRanges.empty() && "expected at least one loop range");
+  assert(loopRanges.size() == tileSizeVals.size() &&
+         "expected as many tile sizes as loop ranges");
+  OpBuilder::InsertionGuard guard(builder);
+  SmallVector<scf::ForOp> loops;
+  offsets.resize(loopRanges.size());
+  sizes.resize(loopRanges.size());
+
+  // The tile size to use (to avoid out of bounds access) is  minimum of
+  // `tileSize` and `ub - iv`, where `iv` is the induction variable
+  // of the tiled loop.
+  AffineExpr s0, s1, d0;
+  bindDims(builder.getContext(), d0);
+  bindSymbols(builder.getContext(), s0, s1);
+  AffineMap minMap = AffineMap::get(1, 2, {s0, s1 - d0}, builder.getContext());
+
+  for (auto loopRange : llvm::enumerate(loopRanges)) {
+    Value offset =
+        getValueOrCreateConstantIndexOp(builder, loc, loopRange.value().offset);
+    Value size =
+        getValueOrCreateConstantIndexOp(builder, loc, loopRange.value().size);
+    // No loops if tile size is zero. Set offset and size to the loop
+    // offset and size.
+    if (matchPattern(tileSizeVals[loopRange.index()], m_Zero())) {
+      offsets[loopRange.index()] = offset;
+      sizes[loopRange.index()] = size;
+      continue;
+    }
+
+    auto loop = builder.create<scf::ForOp>(
+        loc, offset, size, tileSizeVals[loopRange.index()], ValueRange{},
+        [&](OpBuilder &bodyBuilder, Location bodyLoc, Value iv,
+            ValueRange /*iterArgs*/) {
+          bool canAvoidMap = tileDividesIterationDomain(
+              Range{loopRange.value().offset, loopRange.value().size,
+                    tileSizeVals[loopRange.index()]});
+          Value boundedTileSize =
+              (canAvoidMap)
+                  ? tileSizeVals[loopRange.index()]
+                  : builder.create<AffineMinOp>(
+                        bodyLoc, minMap,
+                        ValueRange{iv, tileSizeVals[loopRange.index()], size});
+          sizes[loopRange.index()] = boundedTileSize;
+          builder.create<scf::YieldOp>(loc);
+        });
+    offsets[loopRange.index()] = loop.getInductionVar();
+    loops.push_back(loop);
+    builder.setInsertionPoint(loop.getBody()->getTerminator());
+  }
+  return loops;
+}
+
+/// For a value to be yielded (`yieldedValue`) from within a loop nest `loops`,
+/// construct the destructive update pattern that inserts the yielded
+/// value into a destination tensor provided by `initValue` at offset
+/// `tileOffsets` and size `tileSizes`. For example,
+///
+/// ```mlir
+/// scf.for %iv0 = ... {
+///   %0 = tiled_op
+/// }
+/// ```
+///
+/// is transformed to
+///
+/// ```mlir
+/// scf.for %iv0 = ... iter_args(%arg = %0) {
+///   %1 = tensor.extract_slice %arg
+///   %2 = tiled_op
+///   %3 = tensor.insert_slice %2 into %arg
+///   scf.yield %3
+/// }
+/// ```
+/// TODO: This API can be cleaned up by using `SubsetExtractOpInterface`.
+FailureOr<SmallVector<Value>>
+mlir::yieldTiledValues(RewriterBase &rewriter, ValueRange initValues,
+                       ValueRange yieldedValues,
+                       ArrayRef<SmallVector<OpFoldResult>> tileOffsetsList,
+                       ArrayRef<SmallVector<OpFoldResult>> tileSizesList,
+                       MutableArrayRef<scf::ForOp> loops) {
+  NewYieldValueFn yieldValueFn =
+      [&](OpBuilder &b, Location loc,
+          ArrayRef<BlockArgument> newBBArgs) -> SmallVector<Value> {
+    SmallVector<Value> inserts;
+    for (const auto &yieldedValue : llvm::enumerate(yieldedValues)) {
+      ArrayRef<OpFoldResult> tileOffsets =
+          tileOffsetsList[yieldedValue.index()];
+      ArrayRef<OpFoldResult> tileSizes = tileSizesList[yieldedValue.index()];
+      SmallVector<OpFoldResult> tileStrides(tileOffsets.size(),
+                                            b.getIndexAttr(1));
+      Value insert = b.create<tensor::InsertSliceOp>(
+          loc, yieldedValue.value(), newBBArgs[yieldedValue.index()],
+          tileOffsets, tileSizes, tileStrides);
+      inserts.push_back(insert);
+    }
+    return inserts;
+  };
+
+  SmallVector<scf::ForOp> newLoops =
+      replaceLoopNestWithNewYields(rewriter, loops, initValues, yieldValueFn,
+                                   /*replaceIterOperandsUsesInLoop =*/false);
+  for (const auto &loop : llvm::enumerate(loops)) {
+    rewriter.eraseOp(loop.value());
+    loops[loop.index()] = newLoops[loop.index()];
+  }
+  return llvm::to_vector(llvm::map_range(
+      loops.front().getResults().take_back(yieldedValues.size()),
+      [](OpResult r) -> Value { return r; }));
+}
+
+/// If the tiled operation is destination passing style, update the
+/// slice of the destination used (which refers to the untiled destination)
+/// to use the corresponding region argument of the innermost loop.
+///
+/// ```mlir
+/// %0 =
+/// scf.for %iv0 = ... iter_args(%arg = %0) {
+///   %1 = tensor.extract_slice %0
+///   %2 = tiled_op
+///   %3 = tensor.insert_slice %2 into %arg
+///   scf.yield %3
+/// }
+/// ```
+///
+/// is transformed to
+///
+/// ```mlir
+/// scf.for %iv0 = ... iter_args(%arg = %0) {
+///   %1 = tensor.extract_slice %arg
+///   %2 = tiled_op
+///   %3 = tensor.insert_slice %2 into %arg
+///   scf.yield %3
+/// }
+/// ```
+void mlir::updateDestinationOperandsForTiledOp(
+    OpBuilder &builder, ValueRange tiledOpDestinationValues,
+    ValueRange bbArgsList) {
+  for (const auto &destValue : llvm::enumerate(tiledOpDestinationValues)) {
+    auto sliceOp = destValue.value().getDefiningOp<tensor::ExtractSliceOp>();
+    if (!sliceOp)
+      continue;
+    sliceOp.setOperand(0, bbArgsList[destValue.index()]);
+  }
+}