[mlir] Handle linalg.index correctly in TilingInterface

[ftynse]

Summary

The existing implementation of the TilingInterface for Linalg ops was not
modifying the linalg.index ops contained within other Linalg ops (they need
to be summed up with the values of respective tile loop induction variables),
which led to the interface-based tiling being incorrect for any Linalg op with
index semantics.

In the process, fix the function performing the index offsetting to use the
pattern rewriter API instead of RAUW as it is being called from patterns and
may mess up the internal state of the rewriter. Also rename the function to
clearly catch all uses.

Depends On D129365

[ftynse]

Fix test.

[mravishankar]

Thanks Alex!

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

 /// number of values in `ivs`.
 SmallVector<Value, 4> makeTiledShapes(OpBuilder &builder, Location loc,
                                       LinalgOp linalgOp,
                                       ArrayRef<Value> valuesToTile,
                                       ValueRange ivs, ValueRange tileSizes,
                                       ArrayRef<Value> sizeBounds,
                                       bool omitPartialTileCheck);
 
-/// Add the tile loop induction variables `ivs` to the IndexOp results found in
-/// the body of the `tiledOp` to account for the tile offset.
-void addTileLoopIvsToIndexOpResults(OpBuilder &b, LinalgOp tiledOp,
-                                    ArrayRef<Value> ivs);
+/// Add the specified offsets to any `linalg.index` ops contained in the given
+/// `linalgOp`. The offsets are provided in the same order as iteration space
+/// dimensions. Null offests are assumed to be zero.
+void offsetIndices(OpBuilder &b, LinalgOp linalgOp, ArrayRef<Value> offests);
+void offsetIndices(RewriterBase &b, LinalgOp linalgOp, ArrayRef<Value> offests);
 
 using FusableOpDependencesTy = llvm::MapVector<
     Operation *,
     SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 1>>;
 FusableOpDependencesTy
 findAllFusableDependences(ArrayRef<LinalgOp> ops,
                           const LinalgDependenceGraph &dependenceGraph);
 

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

static LinalgOp fuse(OpBuilder &b, LinalgOp producer,
   }
 
   Operation *clonedOp = producer.clone(b, loc, resultTypes, clonedShapes);
 
   // Shift all IndexOp results by the tile offset.
   SmallVector<Value> allIvs;
   llvm::transform(loopRanges, std::back_inserter(allIvs),
                   [](Range range) { return range.offset; });
-  addTileLoopIvsToIndexOpResults(b, clonedOp, allIvs);
+  offsetIndices(b, clonedOp, allIvs);
 
   return clonedOp;
 }
 
 /// Get the loop range for a dimension `dim` based on the `shapedOperand`. It is
 /// expected to be defined by a subview op or an extract_slice op.
 static Range getRangeFromOperandShape(OpBuilder &b, Location loc,
                                       Value shapedOperand, unsigned dim) {

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

static LinalgOp getTiledProducer(OpBuilder &b, OpResult producerResult,
 
   // Clone the producer using the tiled producer operands.
   TypeRange resultTypes = ValueRange(tiledOperands)
                               .take_back(producerOp.getNumOutputs())
                               .getTypes();
   LinalgOp clonedOp = producerOp.clone(b, loc, resultTypes, tiledOperands);
 
   // Shift all IndexOp results by the tile offset.
-  addTileLoopIvsToIndexOpResults(b, clonedOp, allIvs);
+  offsetIndices(b, clonedOp, allIvs);
 
   return clonedOp;
 }
 
 //===----------------------------------------------------------------------===//
 // TileLoopNest specific helpers.
 //===----------------------------------------------------------------------===//
 

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

std::pair<LinalgOp, LinalgOp> linalg::splitOp(RewriterBase &rewriter,
   LinalgOp second =
       createSplitPart(builder, op, firstResults, splitIterationSpace, dimension,
                       remainingSizes.front(), secondResults, splitPointValue);
 
   // Fixup the linalg.index results in the second part.
   SmallVector<Value> ivAdditions;
   ivAdditions.resize(splitIterationSpace.size());
   ivAdditions[dimension] = splitPointValue;
-  linalg::addTileLoopIvsToIndexOpResults(builder, cast<LinalgOp>(second),
-                                         ivAdditions);
+  linalg::offsetIndices(rewriter, cast<LinalgOp>(second), ivAdditions);
 
   // Replace the original op with the results of the two newly created ops.
   rewriter.replaceOp(op, secondResults);
   return std::make_pair(first, second);
 }

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

void mlir::linalg::transformIndexOps(
     const LoopIndexToRangeIndexMap &loopIndexToRangeIndex) {
   SmallVector<Value> allIvs(op.getNumLoops(), nullptr);
   for (auto &en : enumerate(allIvs)) {
     auto rangeIndex = loopIndexToRangeIndex.find(en.index());
     if (rangeIndex == loopIndexToRangeIndex.end())
       continue;
     en.value() = ivs[rangeIndex->second];
   }
-  addTileLoopIvsToIndexOpResults(b, op, allIvs);
+  offsetIndices(b, op, allIvs);
 }
 
 /// Asserts that the given index-typed value is strictly positive. If the value
 /// is an attribute, asserts at compile time, otherwise emits an assertion
 /// checked at runtime.
 static void emitIsPositiveIndexAssertion(ImplicitLocOpBuilder &b,
                                          OpFoldResult value) {
   if (auto attr = value.dyn_cast<Attribute>()) {

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

getTiledImplementation(Operation *op, OpBuilder &b, ValueRange dest,
                          ArrayRef<OpFoldResult> offsets,
                          ArrayRef<OpFoldResult> sizes,
                          bool tileDestOperands) const {
     // Leave the `sizeBounds` value empty. That is only needed when the `sizes`
     // specified could lead to out of bounds accesses.
     Location loc = op->getLoc();
     LinalgOp linalgOp = cast<LinalgOp>(op);
     SmallVector<Value> valuesToTile = linalgOp.getInputAndOutputOperands();
+    SmallVector<Value> offsetValues =
+        getValueOrCreateConstantIndexOp(b, loc, offsets);
     SmallVector<Value, 4> tiledOperands = makeTiledShapes(
-        b, loc, linalgOp, valuesToTile,
-        getValueOrCreateConstantIndexOp(b, loc, offsets),
+        b, loc, linalgOp, valuesToTile, offsetValues,
         getValueOrCreateConstantIndexOp(b, loc, sizes), {}, true);
 
     SmallVector<Type> resultTensorTypes = llvm::to_vector(llvm::map_range(
         linalgOp.getOutputTensorOperands(), [&](OpOperand *opOperand) {
           return tiledOperands[opOperand->getOperandNumber()].getType();
         }));
 
     Operation *tiledOp =
         linalgOp.clone(b, loc, resultTensorTypes, tiledOperands);
+    offsetIndices(b, cast<LinalgOp>(tiledOp), offsetValues);
 
     return {tiledOp};
   }
 
   // Return the details of the output tile generated by the tiled
   // implementation.
   LogicalResult
   getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,

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

for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
     tiledShapes.push_back(makeTiledShape(b, loc, shapedOp, tileSizes, map, lbs,
                                          sizeBounds, subShapeSizes,
                                          omitPartialTileCheck));
   }
 
   return tiledShapes;
 }
 
-void addTileLoopIvsToIndexOpResults(OpBuilder &b, LinalgOp tiledOp,
-                                    ArrayRef<Value> ivs) {
-  if (tiledOp.hasIndexSemantics()) {
-    for (IndexOp indexOp : tiledOp.getBlock()->getOps<IndexOp>()) {
-      if (ivs[indexOp.dim()] == nullptr)
-        continue;
-      OpBuilder::InsertionGuard guard(b);
-      b.setInsertionPointAfter(indexOp);
-      AffineExpr index, offset;
-      bindDims(b.getContext(), index, offset);
-      AffineApplyOp applyOp = makeComposedAffineApply(
-          b, indexOp.getLoc(), index + offset,
-          ValueRange{indexOp.getResult(), ivs[indexOp.dim()]});
-      indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp);
-    }
+void offsetIndices(OpBuilder &b, LinalgOp linalgOp, ArrayRef<Value> offsets) {
+  IRRewriter rewriter(b);
+  offsetIndices(rewriter, linalgOp, offsets);
+}
+
+void offsetIndices(RewriterBase &b, LinalgOp linalgOp,
+                   ArrayRef<Value> offsets) {
+  if (!linalgOp.hasIndexSemantics())
+    return;
+
+  for (IndexOp indexOp : linalgOp.getBlock()->getOps<IndexOp>()) {
+    if (indexOp.dim() >= offsets.size() || offsets[indexOp.dim()] == nullptr)
+        ValueRange{indexOp.getResult(), offsets[indexOp.dim()]});
+    b.replaceOpWithIf(indexOp, applyOp.getResult(), [&](OpOperand &use) {
+      return use.getOwner() != applyOp;
+    });
   }
 }
 
 /// Get the reassociation maps to fold the result of a extract_slice (or source
 /// 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

mlir/test/Interfaces/TilingInterface/tile-using-interface.mlir

 //  CHECK-DAG:         %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT2]]
 // CHECK-SAME:             [0, 0, 0, 0] [%[[N]], %[[R]], %[[S]], %[[F]]]
 //      CHECK:         %[[CONV_TILE:.+]] = linalg.conv_2d_nhwc_hwcf
 // CHECK-SAME:             dilation = dense<[4, 5]> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64>
 // CHECK-SAME:             ins(%[[INPUT_TILE]], %[[FILTER_TILE]] :
 // CHECK-SAME:             outs(%[[INIT_TILE]] :
 //      CHECK:         tensor.insert_slice %[[CONV_TILE]] into %[[INIT2]]
 // CHECK-SAME:             [0, 0, 0, 0] [%[[N]], %[[R]], %[[S]], %[[F]]]
+
+// -----
+
+// CHECK: #[[$MAP_ADD:.+]] = affine_map<(d0, d1) -> (d0 + d1)>
+
+// CHECK-LABEL: @indexed_semantics
+func.func @indexed_semantics(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  // Check that we correctly amend "linalg.index" results.
+
+  // CHECK: scf.for %[[I0:.+]] = %{{.*}} to %{{.*}} step %{{.*}}
+  // CHECK: scf.for %[[I1:.+]] = %{{.*}} to %{{.*}} step %{{.*}}
+  %0 = linalg.generic {
+    indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                     affine_map<(d0, d1) -> (d0, d1)>],
+    iterator_types = ["parallel", "parallel"]}
+    {__internal_linalg_transform__ = "indexed_semantics"}
+    ins(%arg0: tensor<?x?xf32>)
+    outs(%arg1: tensor<?x?xf32>) {
+  ^bb0(%arg2: f32, %arg3: f32):
+    // CHECK: %[[INDEX0:.+]] = linalg.index 0
+    // CHECK: %[[INDEX0_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX0]], %[[I0]])
+    %1 = linalg.index 0 : index
+    // CHECK: %[[INDEX1:.+]] = linalg.index 1
+    // CHECK: %[[INDEX1_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX1]], %[[I1]])
+    %2 = linalg.index 1 : index
+    // CHECK: arith.addi %[[INDEX0_AMENDED]], %[[INDEX1_AMENDED]]
+    %3 = arith.addi %1, %2 : index
+    %4 = arith.index_cast %3 : index to i64
+    %5 = arith.uitofp %4 : i64 to f32
+    %6 = arith.addf %5, %arg2 : f32
+    linalg.yield %6 : f32
+  } -> (tensor<?x?xf32>)
+  return %0 : tensor<?x?xf32>
+}

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterface.cpp

if (testTiling) {
     addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
         context, {10, 20, 30}, "simple_gemm_memref", patterns);
     // 3. Tiling 3D parallel generic op which implements a transpose
     addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
         context, {10, 0, 20}, "parallel_generic_transpose", patterns);
     // 4. Tiling 2D conv op.
     addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
         context, {0, 0, 0, 0, 10, 20, 30}, "simple_conv", patterns);
+    // 5. Tiling a simple op with `linalg.index` inside.
+    addPatternForTiling<TestTileUsingSCFForOpWithFilter>(
+        context, {10, 20}, "indexed_semantics", patterns);
     return;
   }
   if (testTileConsumerAndFuseProducer) {
     // 1. Tile and fuse of gemm with bias-add operation.
     addPatternForTiling<
         TestTileConsumerAndFuseProducersUsingSCFForOpWithFilter>(
         context, {10, 20}, "fusion", patterns);
     addPatternForTiling<