[MLIR] TilingInterface: Avoid map when tile divides iteration domain

[ftynse]

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

137

Can't we use makeComposedFoldedAffineMin and let it fold the map?

[chelini]

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

137

Thanks for the pointer. Even using makeComposedFoldedAffineMin or makeComposedAffineMin I cannot get rid of the map, but we get better IR with what we have now where the constants are folded into the map.

For example, affine_map<(d0)[s0, s1] -> (20, -d0 + s1)> turns to affine_map<(d0) -> (20, -d0 + 300)>. Current IR:

#map1 = affine_map<(d0) -> (20, -d0 + 300)>

scf.for %arg4 = %c0 to %c300 step %c20 {
   %0 = affine.min #map1(%arg4)
    use(%0)
}

Do we have something to simplify the map here? As a reference also here we do check if we can avoid the map: https://github.com/llvm/llvm-project/blob/bc32896e9f39c1c64528840d78c8a07aad2ad313/mlir/lib/Dialect/Linalg/Utils/Utils.cpp#L866

[ftynse]

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

93

Take either a Range or just three separate variables to avoid magic number-based access below.

100–104

This matching is insufficiently robust, m_Constant is recommended instead. You can also just call getAsOpFoldResult for the stride (offset and size are already available as OpFoldResult at the call side) and dyn_cast those to Attribute.

137

I see, you need to be able to reason about the set of induction variable values. Scrap that.

152

Note that offset and size are created artificially from an OpFoldResult a couple of lines above. Use that instead and you won't do the redundant work of matching constant operations that you have just created from actual constants.

[chelini]

Address comments.

[chelini]

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

100–104

found getConstantIntValue that does what suggested.

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

static bool isPermutation(ArrayRef<unsigned> interchange) {
   }
   return seenVals.size() == interchange.size();
 }
 
 //===----------------------------------------------------------------------===//
 // TileUsingSCFForOp pattern 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,
if (matchPattern(tileSizeVals[loopRange.index()], m_Zero())) {
       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*/) {
-          Value boundedTileSize = builder.create<AffineMinOp>(
+          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;

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

%0:2 = linalg.generic {
       ins(%arg0 : tensor<128x200x300xf32>)
       outs(%init0, %init1 : tensor<128x300x200xf32>, tensor<300x128x200xf32>) {
     ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
       linalg.yield %b0, %b0 : f32, f32
     } -> (tensor<128x300x200xf32>, tensor<300x128x200xf32>)
   return %0#0, %0#1 : tensor<128x300x200xf32>, tensor<300x128x200xf32>
 }
 //  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)>
-//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0, s1] -> (20, -d0 + s1)>
 //      CHECK: func.func @multi_result(
 // CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<128x200x300xf32>)
 //  CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
 //  CHECK-DAG:   %[[C10:.+]] = arith.constant 10 : index
 //  CHECK-DAG:   %[[C20:.+]] = arith.constant 20 : index
 //  CHECK-DAG:   %[[C128:.+]] = arith.constant 128 : index
 //  CHECK-DAG:   %[[C300:.+]] = arith.constant 300 : index
 //  CHECK-DAG:   %[[INIT0:.+]] = linalg.init_tensor [128, 300, 200]
 //  CHECK-DAG:   %[[INIT1:.+]] = linalg.init_tensor [300, 128, 200]
 //      CHECK:   %[[OUTER:[a-zA-Z0-9]+]]:2 = scf.for %[[IV0:[a-zA-Z0-9]+]] = %[[C0]] to %[[C128]] step %[[C10]]
 // CHECK-SAME:       iter_args(%[[ARG1:[a-zA-Z0-9]+]] = %[[INIT0]], %[[ARG2:[a-zA-Z0-9]+]] = %[[INIT1]])
 //      CHECK:     %[[TS_Y:.+]] = affine.min #[[MAP0]](%[[IV0]])[%[[C10]], %[[C128]]]
 //      CHECK:     %[[INNER:[a-zA-Z0-9]+]]:2 = scf.for %[[IV1:[a-zA-Z0-9]+]] = %[[C0]] to %[[C300]] step %[[C20]]
 // CHECK-SAME:         iter_args(%[[ARG3:[a-zA-Z0-9]+]] = %[[ARG1]], %[[ARG4:[a-zA-Z0-9]+]] = %[[ARG2]])
-//      CHECK:       %[[TS_X:.+]] = affine.min #[[MAP1]](%[[IV1]])[%[[C20]], %[[C300]]]
 //  CHECK-DAG:       %[[ARG_TILE:.+]] = tensor.extract_slice %[[ARG0]]
-// CHECK-SAME:           [%[[IV0]], 0, %[[IV1]]] [%[[TS_Y]], 200, %[[TS_X]]] [1, 1, 1]
+// CHECK-SAME:           [%[[IV0]], 0, %[[IV1]]] [%[[TS_Y]], 200, 20] [1, 1, 1]
 //  CHECK-DAG:       %[[INIT0_TILE:.+]] = tensor.extract_slice %[[ARG3]]
-// CHECK-SAME:           [%[[IV0]], %[[IV1]], 0] [%[[TS_Y]], %[[TS_X]], 200] [1, 1, 1]
+// CHECK-SAME:           [%[[IV0]], %[[IV1]], 0] [%[[TS_Y]], 20, 200] [1, 1, 1]
 //  CHECK-DAG:       %[[INIT1_TILE:.+]] = tensor.extract_slice %[[ARG4]]
-// CHECK-SAME:           [%[[IV1]], %[[IV0]], 0] [%[[TS_X]], %[[TS_Y]], 200] [1, 1, 1]
+// CHECK-SAME:           [%[[IV1]], %[[IV0]], 0] [20, %[[TS_Y]], 200] [1, 1, 1]
 //      CHECK:       %[[RESULT_TILE:.+]]:2 = linalg.generic
 // CHECK-SAME:           ins(%[[ARG_TILE]] :
 // CHECK-SAME:           outs(%[[INIT0_TILE]], %[[INIT1_TILE]] :
 //      CHECK:       %[[UPDATE0:.+]] = tensor.insert_slice %[[RESULT_TILE]]#0 into %[[ARG3]]
-// CHECK-SAME:           [%[[IV0]], %[[IV1]], 0] [%[[TS_Y]], %[[TS_X]], 200] [1, 1, 1]
+// CHECK-SAME:           [%[[IV0]], %[[IV1]], 0] [%[[TS_Y]], 20, 200] [1, 1, 1]
 //      CHECK:       %[[UPDATE1:.+]] = tensor.insert_slice %[[RESULT_TILE]]#1 into %[[ARG4]]
-// CHECK-SAME:           [%[[IV1]], %[[IV0]], 0] [%[[TS_X]], %[[TS_Y]], 200] [1, 1, 1]
+// CHECK-SAME:           [%[[IV1]], %[[IV0]], 0] [20, %[[TS_Y]], 200] [1, 1, 1]
 //      CHECK:       scf.yield %[[UPDATE0]], %[[UPDATE1]]
 //      CHECK:     scf.yield %[[INNER]]#0, %[[INNER]]#1
 //      CHECK:   return %[[OUTER]]#0, %[[OUTER]]#1
 
 // -----
 
 func.func @conv2D(%arg0 : tensor<?x?x?x?xf32>, %arg1 : tensor<?x?x?x?xf32>,
     %arg2 : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {