[mlir][linalg] Fold ExtractSliceOps during tiling.

[gysit]

Summary

Add the makeComposedExtractSliceOp method that creates an ExtractSliceOp and folds chains of ExtractSliceOps by computing the sum of their offsets and by multiplying their strides.

[gysit]

Fix windows build.

[mravishankar]

Just a few comments.

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

206

I have written this function multiple times. Need a common place for it. I think its complicated since it requires standard dialect so any place this is put needs to also add dependency on standard dialect.

219

Not sure doing this recursively in a loop is a good idea. Maybe just do one step and allow the pattern rewriter to apply it multiple times?

221

Why? This is just folding the offsets and the strides. The type should still remain the same right. So you should be able to handle this?

228

I am not convinced that you can just add the offsets. I think you can do that only if the strides are 1 (though I havent given this deep thought, just a first impression).

I would just simplify all of this and check for all strides in the producer and consumer to be unit strides and bail if they arent. I dont think we have a case with strides not being 1 anyway.

Meta point: We should just drop strides from subviews. Do we ever use non-unit strides?

[gysit]

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

206

I tried to reuse an existing implementation but all of them introduce strange dependencies. Also a standard dialect implementation probably requires more functionality in the sense that we would have to support non integer types etc.

221

The idea is to make sure the producer is not rank reducing. It it is there would be a type mismatch.

228

Good point! I think you are right regarding the strides and adding offsets. The two offsets can potentially have different strides which cannot be handled properly in one extract slice. I will check for stride one then.

[springerm]

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

206

Just fyi. There is getAsOpFoldResult etc. in StaticValueUtils, which are related to this. I ran into the same issue with "dependency on std" when I added these, which is why I couldn't add this function.

Anyway... I don't know how to solve this in a nice way either.

[gysit]

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

206

StandardOps/Utils/Utils.h may be an option? But then the idea seems to be to remove std...

[gysit]

Address comments.

[gysit]

Cleaner version of OpFoldResult to Value helper.

[nicolasvasilache]

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

203

can you early exit and avoid increasing nesting logic?

209

Just checking equality of the strides of the 2 ops would be enough here.
It would also capture rank preserving.

Re. "do we need something else than stride 1": in practice today no but things like dilation in convolutions are one of the envisioned use cases.

223

Could you provide a version of makeComposedAffineApply that works with OpFoldResult and simplify the caller?

[gysit]

Address comments.

[gysit]

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

209

After some more thought I think the relevant information are the producer strides? If they are one then both offsets are expressed using the same strides.

223

After some more though I think my last version was wrong as well. If I am not mistaken the best solution for now is to check the producer strides are one. I think this is a necessary condition to guarantee the producer and consumer offsets are expressed in the same stride?

In theory we could also handle other producer strides by multiplying the consumer offset by the producer stride before adding the offsets for the folding. Or am I wrong again here?

[nicolasvasilache]

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

209

Ah yes my apologies .. strides in subX operations compose (i.e. the op can strictly subsample its input).

Resulting strides should be the product of strides.
Resulting sizes should be the consumer sizes.
Result offset should consumer_offset + producer_offset * consumer strides.

[nicolasvasilache]

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

209

Sigh .. I meant:

Result offset should be producer_offset + consumer_offset * producer_strides.

Here is the code that implements the lowering to LLVM (replace producer by "sourceMemRef"):
https://sourcegraph.com/github.com/llvm/llvm-project/-/blob/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp?L1307

Now as Mahesh pointed out, this is not well tested on actual use cases so maybe failing if strides are != 1 is the most pragmatic thing to do right now ..

[gysit]

Address comment.
In particluar:

• Support unit-strides only.

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

 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
 
 /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp
 /// is a constant then return a new value set to the smallest such constant.
 /// If `size` comes from a ConstantOp, return the constant.
 /// Otherwise return nullptr.
 IntegerAttr getSmallestBoundingIndex(Value size);
 
+/// Create an ExtractSliceOp and, if `source` is defined by an ExtractSliceOp
+/// chain, fold it by adding the offsets and multiplying the strides.
+///
+/// Example:
+/// ```
+///   %0 = tensor.extract_slice %arg0[3, 0][3, 32][1, 2] : tensor<64x64xf32> to
+///                                                        tensor<3x32xf32>
+///   %1 = tensor.extract_slice %0[0, 0][3, 4][1, 3] : tensor<3x32xf32> to
+///                                                    tensor<3x4xf32>
+/// ```
+/// folds into:
+/// ```
+///   %1 = tensor.extract_slice %arg0[3, 0] [3, 4] [1, 6] : tensor<64x64xf32> to
+///                                                         tensor<3x4xf32>
+/// ```
+tensor::ExtractSliceOp makeComposedExtractSliceOp(
+    OpBuilder &b, Location loc, Value source, ArrayRef<OpFoldResult> offsets,
+    ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides);
+
 //===----------------------------------------------------------------------===//
 // Fusion utilities
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the specific `producer` is the last write to exactly the
 /// whole `consumedView`. This checks structural dominance, that the dependence
 /// is a RAW without any interleaved write to any piece of `consumedView`.
 bool isProducerLastWriteOfView(const LinalgDependenceGraph &graph,

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

if (auto constOp = dimOp.index().getDefiningOp<ConstantOp>()) {
       }
     }
   }
   if (boundingConst && *boundingConst >= 0)
     return Builder(size.getContext()).getIndexAttr(*boundingConst);
   return nullptr;
 }
 
+tensor::ExtractSliceOp makeComposedExtractSliceOp(
+    OpBuilder &b, Location loc, Value source, ArrayRef<OpFoldResult> offsets,
+    ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) {
+  AffineExpr dim1, dim2;
+  AffineExpr sym1, sym2;
+  bindDims(b.getContext(), dim1, dim2);
+  bindSymbols(b.getContext(), sym1, sym2);
+
+  // Helper function to convert fold results into values.
+  auto getOpFoldResultAsValue = [&](OpFoldResult ofr) -> Value {
+    if (auto result = ofr.dyn_cast<Value>())
+      return result;
+    assert(ofr.dyn_cast<Attribute>().isa<IntegerAttr>() &&
+           "expect only integer op fold results");
+    APInt apInt = ofr.dyn_cast<Attribute>().cast<IntegerAttr>().getValue();
+    return b.create<ConstantIndexOp>(loc, apInt.getSExtValue());
+  };
+
+  // Accumulate the `source` offsets and strides.
+  SmallVector<OpFoldResult> foldedOffsets(offsets.begin(), offsets.end());
+  SmallVector<OpFoldResult> foldedStrides(strides.begin(), strides.end());
+  assert(source && "expect source to be nonezero");
+  while (auto producerOp = source.getDefiningOp<tensor::ExtractSliceOp>()) {
+    // Ensure the producer is not rank reducing.
+    if (producerOp.getSourceType().getRank() != offsets.size())
+      break;
+    // Add the producer offsets.
+    for (auto en : enumerate(producerOp.getMixedOffsets())) {
+      SmallVector<Value> offsetValues = {
+          getOpFoldResultAsValue(en.value()),
+          getOpFoldResultAsValue(foldedOffsets[en.index()])};
+      foldedOffsets[en.index()] =
+          makeComposedAffineApply(b, loc, dim1 + dim2, offsetValues)
+              .getResult();
+    }
+    // Multiply the producer strides.
+    for (auto en : enumerate(producerOp.getMixedStrides())) {
+      SmallVector<Value> strideValues = {
+          getOpFoldResultAsValue(en.value()),
+          getOpFoldResultAsValue(foldedStrides[en.index()])};
+      foldedStrides[en.index()] =
+          makeComposedAffineApply(b, loc, sym1 * sym2, strideValues)
+              .getResult();
+    }
+    source = producerOp.source();
+  }
+
+  // Create the ExtractSliceOp using the folded offsets and strides.
+  return b.create<tensor::ExtractSliceOp>(loc, source, foldedOffsets, sizes,
+                                          foldedStrides);
+}
+
 /// Specialization to build an scf "for" nest.
 template <>
 void GenerateLoopNest<scf::ForOp>::doit(
     OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
     ArrayRef<Attribute> iteratorTypes,
     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
                                   ValueRange)>
         bodyBuilderFn,
for (unsigned r = 0; r < rank; ++r) {
 
     sizes.push_back(size);
     LLVM_DEBUG(llvm::dbgs()
                << "makeTiledShape: new offset: " << offset << "\n");
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: new size: " << size << "\n");
     strides.push_back(builder.getIndexAttr(1));
   }
 
-  Operation *sliceOp = shapedType.isa<MemRefType>()
-                           ? builder
-                                 .create<memref::SubViewOp>(
-                                     loc, valueToTile, offsets, sizes, strides)
-                                 .getOperation()
-                           : builder
-                                 .create<tensor::ExtractSliceOp>(
-                                     loc, valueToTile, offsets, sizes, strides)
-                                 .getOperation();
+  Operation *sliceOp =
+      shapedType.isa<MemRefType>()
+          ? builder.create<memref::SubViewOp>(loc, valueToTile, offsets, sizes,
+                                              strides)
+          : makeComposedExtractSliceOp(builder, loc, valueToTile, offsets,
+                                       sizes, strides);
+
   return sliceOp->getResult(0);
 }
 
 SmallVector<Value> computeTileOffsets(OpBuilder &b, Location loc,
                                       ValueRange ivs, ValueRange tileSizes) {
   SmallVector<Value> offsets;
   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n");

mlir/test/Dialect/Linalg/tile-tensors.mlir

 
 // TLOOP:     %{{.*}} = linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) =
 // TLOOP-SAME: (%[[C0]], %[[C0]], %[[C0]])
 // TLOOP-SAME: to (%[[ARG_0_X]], %[[ARG_0_Y]], %[[ARG_0_Z]])
 // TLOOP-SAME: step (%[[C2]], %[[C3]], %[[C4]])
 // TLOOP-SAME: ins (%{{.*}} = %[[ARG_0]]: [[TY]], %{{.*}} = %[[ARG_1]]: [[TY]])
 // TLOOP-SAME: outs (%{{.*}} = %[[INIT]]: [[TY]])
 // TLOOP-SAME: distribution["block_x", "block_y", "none"] {
+
+// -----
+
+//  CHECK-DAG:  #[[MAP0:.*]] = affine_map<(d0)[s0] -> (2, -d0 + s0)>
+//  CHECK-DAG:  #[[MAP1:.*]] = affine_map<(d0) -> (d0 + 3)>
+//  CHECK-DAG:  #[[MAP2:.*]] = affine_map<(d0) -> (d0 + 4)>
+
+//      CHECK:  fold_extract_slice
+// CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]*]]: tensor<?x128xf32>
+// CHECK-SAME:    %[[ARG1:[0-9a-zA-Z]*]]: tensor<?x42xf32>
+func @fold_extract_slice(
+  %arg0 : tensor<?x128xf32>, %arg1 : tensor<?x42xf32>, %arg2 : tensor<?x42x?xf32>) -> tensor<?x42xf32> {
+
+  //      CHECK:    %[[C0:.*]] = constant 0
+  %c0 = constant 0 : index
+
+  //      CHECK:    %[[DIM:.*]] = tensor.dim %[[ARG1]], %[[C0]]
+  %0 = tensor.dim %arg1, %c0 : tensor<?x42xf32>
+  %1 = tensor.extract_slice %arg0[3, 4] [%0, 42] [%0, 2] : tensor<?x128xf32> to tensor<?x42xf32>
+
+  //      CHECK:    scf.for %[[IV0:[0-9a-zA-Z]*]] =
+  //      CHECK:      scf.for %[[IV1:[0-9a-zA-Z]*]] =
+
+  // Fold the existing extract slice op into the one created by the tiling.
+  //      CHECK:        %[[SIZE0:.*]] = affine.min #[[MAP0]](%[[IV0]])[%[[DIM]]
+  //      CHECK:        %[[OFF0:.*]] = affine.apply #[[MAP1]](%[[IV0]]
+  //      CHECK:        %[[OFF1:.*]] = affine.apply #[[MAP2]](%[[IV1]]
+  //      CHECK:        %[[T0:.*]] = tensor.extract_slice %[[ARG0]]
+  // CHECK-SAME:                                          %[[OFF0]], %[[OFF1]]
+  // CHECK-SAME:                                          %[[SIZE0]], 3
+  // CHECK-SAME:                                          %[[DIM]], 2
+  //      CHECK:        {{.*}} = linalg.generic {{.*}} ins(%[[T0]]
+  %2 = linalg.generic
+    {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>,
+                      affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+                      affine_map<(d0, d1, d2) -> (d0, d1)>],
+     iterator_types = ["parallel", "parallel", "parallel"]}
+    ins(%1, %arg2 : tensor<?x42xf32>, tensor<?x42x?xf32>)
+    outs(%arg1 : tensor<?x42xf32>) {
+    ^bb0(%arg3 : f32, %arg4: f32, %arg5: f32):
+      %5 = addf %arg3, %arg5 : f32
+      linalg.yield %5 : f32
+    } -> tensor<?x42xf32>
+  return %2 : tensor<?x42xf32>
+}
+