[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

207

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.

220

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?

222

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?

229

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

207

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.

222

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

229

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

207

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

207

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

204

can you early exit and avoid increasing nesting logic?

210

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.

224

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

210

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.

224

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

210

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

210

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,
+/// fold it by adding the offsets.
+///
+/// Example:
+/// ```
+///   %0 = tensor.extract_slice %arg0[3, 4][3, 32][1, 1] : tensor<64x64xf32> to
+///                                                        tensor<3x32xf32>
+///   %1 = tensor.extract_slice %0[0, 5][3, 4][1, 1] : tensor<3x32xf32> to
+///                                                    tensor<3x4xf32>
+/// ```
+/// folds into:
+/// ```
+///   %1 = tensor.extract_slice %arg0[3, 9][3, 4][1, 1] : 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/include/mlir/Dialect/StandardOps/Utils/Utils.h

auto newOp =
                                 mixedOffsets, mixedSizes, mixedStrides);
     CastOpFunc func;
     func(rewriter, op, newOp);
 
     return success();
   }
 };
 
+/// Converts an OpFoldResult to a Value. Returns the fold result if it casts to
+/// a Value or creates a ConstantIndexOp if it casts to an IntegerAttribute.
+/// Other attribute types are not supported.
+Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
+                                      OpFoldResult ofr);
+
 /// Helper struct to build simple arithmetic quantities with minimal type
 /// inference support.
 struct ArithBuilder {
   ArithBuilder(OpBuilder &b, Location loc) : b(b), loc(loc) {}
 
   Value _and(Value lhs, Value rhs);
   Value add(Value lhs, Value rhs);
   Value mul(Value lhs, Value rhs);

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

 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/LoopUtils.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "linalg-utils"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::scf;
 
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) {
+  assert(source && "expect source to be nonzero");
+
+  // Do not fold if the producer is not an ExtractSliceOp.
+  auto producerOp = source.getDefiningOp<tensor::ExtractSliceOp>();
+  if (!producerOp)
+    return b.create<tensor::ExtractSliceOp>(loc, source, offsets, sizes,
+                                            strides);
+
+  // Do not fold if the producer is rank reducing or if there are any non-unit
+  // strides. Supporting non-unit strides complicates the offset computation
+  // since the consumer offsets need to be multiplied by the producer strides.
+  // TODO: support non-unit strides once there are use cases.
+  SmallVector<OpFoldResult> allStrides = producerOp.getMixedStrides();
+  allStrides.append(strides.begin(), strides.end());
+  bool hasNonUnitStride = any_of(allStrides, [](OpFoldResult ofr) {
+    return getConstantIntValue(ofr) != static_cast<int64_t>(1);
+  });
+  if (hasNonUnitStride ||
+      producerOp.getSourceType().getRank() !=
+          producerOp.getResult().getType().cast<ShapedType>().getRank())
+    return b.create<tensor::ExtractSliceOp>(loc, source, offsets, sizes,
+                                            strides);
+
+  // Fold the producer by adding the offests and extracting the slice directly
+  // from the producer source tensor.
+  SmallVector<OpFoldResult> foldedOffsets(offsets.begin(), offsets.end());
+  AffineExpr dim1, dim2;
+  bindDims(b.getContext(), dim1, dim2);
+  for (auto en : enumerate(producerOp.getMixedOffsets())) {
+    SmallVector<Value> offsetValues = {
+        getValueOrCreateConstantIndexOp(b, loc, foldedOffsets[en.index()]),
+        getValueOrCreateConstantIndexOp(b, loc, en.value())};
+    foldedOffsets[en.index()] =
+        makeComposedAffineApply(b, loc, dim1 + dim2, offsetValues).getResult();
+  }
+  return b.create<tensor::ExtractSliceOp>(loc, producerOp.source(),
+                                          foldedOffsets, sizes, strides);
+}
+
 /// 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();
+  auto *sliceOp = TypeSwitch<ShapedType, Operation *>(shapedType)
+                      .Case([&](MemRefType) {
+                        return builder.create<memref::SubViewOp>(
+                            loc, valueToTile, offsets, sizes, strides);
+                      })
+                      .Case([&](RankedTensorType) {
+                        return makeComposedExtractSliceOp(
+                            builder, loc, valueToTile, offsets, sizes, strides);
+                      })
+                      .Default([](ShapedType) -> Operation * {
+                        llvm_unreachable("Unexpected shaped type");
+                      });
   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/lib/Dialect/StandardOps/Utils/Utils.cpp

void mlir::getPositionsOfShapeOne(
   for (unsigned pos = 0, e = shape.size(); pos < e && rank > 0; ++pos) {
     if (shape[pos] == 1) {
       dimsToProject.insert(pos);
       --rank;
     }
   }
 }
 
+Value mlir::getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
+                                            OpFoldResult ofr) {
+  if (auto value = ofr.dyn_cast<Value>())
+    return value;
+  auto attr = ofr.dyn_cast<Attribute>().dyn_cast<IntegerAttr>();
+  assert(attr && "expect the op fold result casts to an integer attribute");
+  return b.create<ConstantIndexOp>(loc, attr.getValue().getSExtValue());
+}
+
 Value ArithBuilder::_and(Value lhs, Value rhs) {
   return b.create<AndOp>(loc, lhs, rhs);
 }
 Value ArithBuilder::add(Value lhs, Value rhs) {
   if (lhs.getType().isa<IntegerType>())
     return b.create<AddIOp>(loc, lhs, rhs);
   return b.create<AddFOp>(loc, lhs, rhs);
 }

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] [1, 1] : 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:                                          1, 1
+  //      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>
+}
+