[mlir][Linalg] NFC: Verify tiling on linalg.generic operation on tensors.

[mravishankar]

Summary

With the recent changes to linalg on tensor semantics, the tiling
operations works out-of-the-box for generic operations. Add a test to
verify that and some minor refactoring.

[nicolasvasilache]

It's nice that we actually don't need to change op semantics to support "init_tensor form" and that all semantics can be dealt with just in tiling.
However we need a better abstraction and documentation to avoid leaking the op + init or result tensor implementation detail into the tiling implementation.

A simple

struct {
  LinalgOp op;
  Value result; // only if op has no init_tensor
};

with a proper API and assertions should be enough to separate these concerns.

The harder part will prob. be finding good names for this :)

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

47 ↗

(On Diff #311064)

initialization

48 ↗

(On Diff #311064)

Now that I see you don't need to specify dims because the op must fold away before bufferization,
I am not worried about "the assumption that entire tensor is overwritten", this is the semantics of linalg ops and this is preserved by tiling.
So I would drop the "Proceed with caution" sentence and instead emphasize that this is an op that must fold away before bufferization.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

68 ↗

(On Diff #311064)

This seems fishy to me because it changes the semantics of all ops.

Shouldn't we instead drop the extra dimensions from the getShapesToLoopsMap as discussed in https://reviews.llvm.org/D93076 ?
As I mentioned there, this was written in the context of buffers where the output buffers are part of the shaped operands.
So this needs a little adjustment but I would claim in the other direction.

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

378

Please isolate in a helper function and document it:
"tiling on tensors requires tensors to write form"
"don't change the op semantics to become init_tensors but instead but create result tensors when needed"
"this op must fold away and does not require size operands even in the dynamic case"
etc

409–410

This doc assumed we would evolve the op semantics to allow all ops to convert to an "init_tensor form".
The direction you are pushing avoids that change and it need to be clearly documented.

409–410

this would benefit from a helper function + documentation that mentions the 2 cases:

1. op with init tensor
2. op without init tensor but with a companion result tensor for the purpose of tiling.

414

This is now harder to follow and we will trip over in the future.

I suggest a helper struct that captures the op and potentially the result tensor.
That struct hides the implementation detail for the purpose of tiling and would be queried as if the op was in init_tensor form.
With proper documentation it will make it much easier to understand in 3 months.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

48 ↗

(On Diff #311064)

There are still issues with this op definition I am not entirely clear about. I sent this out for review to start conversion about this. Here are issues that i think need to be solved

1. If it is upto bufferization to make sure that the init_tensor -> scf.for pattern should result in the "init_tensor being folded away" then I would expect it would be better to make sure there is one init_tensor for every tiling loop generated, i.e. init_tensor ops are not CSEed if there are multiple tiling loops. In that case we probably want to add some dummy side-effect for init_tensor
2. I actually dont think having an initial value here is doing anything. Should we drop the "initialization value" above. The reason why I had this was because leaving it undef is a whole other can of worms, but that seems to actually match what is needed here. The proceed with caution was to say that the initialization value is not load-bearing at all. It is overwritten after tiling loops are executed.

My preference right now is to

• make this a side-effecting operation that cannot be CSE-ed.
• drop the initialization value
• Assume that the bufferization is going to fold the op away

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

68 ↗

(On Diff #311064)

Agreed. WIll send out a different patch for that.

[mravishankar]

Address comments and rebase

[mravishankar]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

68 ↗

(On Diff #311064)

Moved the code to https://reviews.llvm.org/D93180 .

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

378

Good Idea. Made a helper struct TiledOp that encompasses all the logic. It does make things cleaner.

[hanchung]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1774–1778 ↗

(On Diff #311455)

nit: drop braces.

https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

E.g.,

// This should also omit braces.  The `for` loop contains only a single statement,
// so it shouldn't have braces.  The `if` also only contains a single simple
// statement (the for loop), so it also should omit braces.
if (isa<FunctionDecl>(D))
  for (auto *A : D.attrs())

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

314

s/a/an

317

s/a/an

320–342

I would add braces between methods and vars because it helps readability. E.g.,

struct TiledOp {
  TiledOp(OpBuilder &b, LinalgOp op) : op(op) {
     ...
   }

  SmallVector<Value, 4> getTiledOperands() {
     ...
   }

  LinalgOp op;
  SmallVector<Value, 1> initTensors;
};

mlir/test/Dialect/Linalg/roundtrip.mlir

722–727 ↗

(On Diff #311455)

We typically put this to invalid.mlir, lets move to there and remove adding -verify-diagnostics to this file.

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

46

remove the blank line?

[mravishankar]

Updating this to use init_tensor operation, but this doesnt work. It
creates a use-after-def issue. This should be resolved when
https://llvm.discourse.group/t/linalg-and-shapes/2421 is addressed.

[mravishankar]

Rebase over recent changes to linalg on tensors

[mravishankar]

Rebase

[hanchung]

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

575

Since you also add IndexedGenericOp, could we add a test for it?

[nicolasvasilache]

+1 for the extra test.

[mravishankar]

Rebase and add extra test

mlir/include/mlir/IR/AffineMap.h

 ///
 /// Returns the map:
 ///
 /// ```mlir
 ///     (i, j, k) -> (i, k, k, j, i, j)
 /// ```
 AffineMap concatAffineMaps(ArrayRef<AffineMap> maps);
 
+/// Returns the map that results from projecting out the dimensions specified in
+/// `projectedDimensions`. The projected dimensions are set to 0.
+///
+/// Example:
+/// 1) map                  : affine_map<(d0, d1, d2) -> (d0, d1)>
+///    projected_dimensions : {2}
+///    result               : affine_map<(d0, d1) -> (d0, d1)>
+///
+/// 2) map                  : affine_map<(d0, d1) -> (d0 + d1)>
+///    projected_dimensions : {1}
+///    result               : affine_map<(d0) -> (d0)>
+///
+/// 3) map                  : affine_map<(d0, d1, d2) -> (d0, d1)>
+///    projected_dimensions : {1}
+///    result               : affine_map<(d0, d1) -> (d0, 0)>
 AffineMap getProjectedMap(AffineMap map,
                           ArrayRef<unsigned> projectedDimensions);
 
 inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
   map.print(os);
   return os;
 }
 } // end namespace mlir

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

static bool isTiled(AffineMap map, ValueRange tileSizes) {
   for (unsigned r = 0; r < map.getNumResults(); ++r)
     if (isTiled(map.getResult(r), tileSizes))
       return true;
   return false;
 }
 
 static SmallVector<Value, 4>
 makeTiledShapes(OpBuilder &b, Location loc, LinalgOp linalgOp,
-                ValueRange operands, AffineMap map, ValueRange ivs,
+                ArrayRef<Value> tiledOperands, AffineMap map, ValueRange ivs,
                 ValueRange tileSizes, ValueRange allShapeSizes) {
-  assert(operands.size() == linalgOp.getShapedOperands().size());
   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
                            [](Value v) { return !isZero(v); })) &&
          "expected as many ivs as non-zero sizes");
 
   using namespace edsc::op;
 
   auto shapeSizes = applyMapToValues(b, loc, map, allShapeSizes);
   // Construct (potentially temporary) mins and maxes on which to apply maps
   // that define tile subshapes.
   SmallVector<Value, 8> lbs, subShapeSizes;
   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
     bool isTiled = !isZero(tileSizes[idx]);
     lbs.push_back(isTiled ? ivs[idxIvs++] : (Value)std_constant_index(0));
     // Before composing, we need to make range a closed interval.
     Value size = isTiled ? tileSizes[idx] : shapeSizes[idx];
     subShapeSizes.push_back(size - std_constant_index(1));
   }
 
-  auto *op = linalgOp.getOperation();
-
   SmallVector<Value, 4> res;
-  res.reserve(op->getNumOperands());
-  for (auto en : llvm::enumerate(operands)) {
+  res.reserve(tiledOperands.size());
+  for (auto en : llvm::enumerate(tiledOperands)) {
     Value shapedOp = en.value();
     ShapedType shapedType = shapedOp.getType().cast<ShapedType>();
     unsigned rank = shapedType.getRank();
     AffineMap map = linalgOp.getIndexingMap(en.index());
     // If the shape is not tiled, we can use it as is.
     if (!isTiled(map, tileSizes)) {
       res.push_back(shapedOp);
       continue;
for (auto en : llvm::enumerate(tiledOperands)) {
     else
       res.push_back(
           b.create<SubTensorOp>(loc, shapedOp, offsets, sizes, strides));
   }
 
   return res;
 }
 
 template <typename LoopTy>
 static Optional<TiledLinalgOp>
 tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ValueRange tileSizes,
                  const LinalgTilingOptions &options) {
   auto nLoops = op.getNumLoops();
   // Initial tile sizes may be too big, only take the first nLoops.
   tileSizes = tileSizes.take_front(nLoops);
 
   if (llvm::all_of(tileSizes, isZero))
     return llvm::None;
 
   if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) {
     // For conv op only support tiling along batch dimension (which is the first
     // loop).
     if (convOp.padding() && !llvm::all_of(tileSizes.drop_front(), isZero))
       return llvm::None;
   }
 
   // 1. Build the tiled loop ranges.
   auto allShapeSizes = op.createFlatListOfOperandDims(b, op.getLoc());
   AffineMap shapeSizesToLoopsMap = op.getShapesToLoopsMap();
   if (!shapeSizesToLoopsMap)
     return llvm::None;
 
   SmallVector<Range, 4> loopRanges;
   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
   std::tie(loopRanges, loopIndexToRangeIndex) = makeTiledLoopRanges(
       b, op.getLoc(), shapeSizesToLoopsMap, allShapeSizes, tileSizes);
+
   SmallVector<Attribute, 4> iteratorTypes;
   for (auto attr :
        enumerate(op.iterator_types().cast<ArrayAttr>().getValue())) {
     if (loopIndexToRangeIndex.count(attr.index()))
       iteratorTypes.push_back(attr.value());
   }
   // If interchangeVector is empty, use the identity. Build the permutation map
   // otherwise.
if (!options.interchangeVector.empty()) {
     applyPermutationToVector(iteratorTypes, interchangeVector);
   }
 
   // 2. Create the tiled loops.
   LinalgOp res = op;
   SmallVector<Value, 4> ivs, tensorResults;
   auto outputTensors = op.getOutputTensors();
   GenerateLoopNest<LoopTy>::doit(
       loopRanges, /*iterArgInitValues*/ outputTensors, iteratorTypes,
       [&](ValueRange localIvs, ValueRange iterArgs) -> scf::ValueVector {
         auto &b = ScopedContext::getBuilderRef();
         auto loc = ScopedContext::getLocation();
         ivs.assign(localIvs.begin(), localIvs.end());
 
         // When an `interchangeVector` is present, it has been applied to the
         // loop ranges and the iterator types. Apply its inverse to the
         // resulting loop `ivs` to match the op definition.
GenerateLoopNest<LoopTy>::doit(
         operands.append(outputBuffers.begin(), outputBuffers.end());
         operands.append(iterArgs.begin(), iterArgs.end());
         SmallVector<Value, 4> tiledOperands =
             makeTiledShapes(b, loc, op, operands, shapeSizesToLoopsMap,
                             interchangedIvs, tileSizes, allShapeSizes);
         auto nonShapedOperands = op.getAssumedNonShapedOperands();
         tiledOperands.append(nonShapedOperands.begin(),
                              nonShapedOperands.end());
 
         // TODO: use an interface/adaptor to avoid leaking position in
         // `tiledOperands`.
         SmallVector<Type, 4> resultTensorTypes;
         for (OpOperand *opOperand : op.getOutputTensorsOpOperands())
           resultTensorTypes.push_back(
               tiledOperands[opOperand->getOperandNumber()].getType());
 
         res = op.clone(b, loc, resultTensorTypes, tiledOperands);
 
         // Insert a subtensor_insert for each output tensor.
         unsigned resultIdx = 0;
         for (OpOperand *opOperand : op.getOutputTensorsOpOperands()) {
           // TODO: use an interface/adaptor to avoid leaking position in

 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >::insert(patterns, ctx);
 }
 
 /// Populate the given list with patterns that apply Linalg tiling.
 static void insertTilingPatterns(OwningRewritePatternList &patterns,
                                  const LinalgTilingOptions &options,
                                  MLIRContext *ctx) {
-  RewritePatternList<
+  RewritePatternList<GenericOp, IndexedGenericOp,
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
                      >::insert(patterns, options, ctx);
 }
 
 static void applyTilingToLoopPatterns(LinalgTilingLoopType loopType,
                                       FuncOp funcOp,
                                       ArrayRef<int64_t> tileSizes) {
   auto options =
       LinalgTilingOptions().setTileSizes(tileSizes).setLoopType(loopType);
   MLIRContext *ctx = funcOp.getContext();

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

-// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,4" | FileCheck %s
+// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,4" -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @matmul_tensors(
 // CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<?x?xf32>
 // CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<?x?xf32>
 // CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
 func @matmul_tensors(
   %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
     -> tensor<?x?xf32> {

 //      CHECK:   scf.yield %[[TD1]] : tensor<?x?xf32>
   %0 = linalg.matmul  ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
                      outs(%arg2: tensor<?x?xf32>)
     -> tensor<?x?xf32>
 
 //      CHECK: return %[[TD0]] : tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
+
+// -----
+
+func @generic_op_tensors(
+  %arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %0 = dim %arg0, %c0 : tensor<?x?x?xf32>
+  %1 = dim %arg0, %c1 : tensor<?x?x?xf32>
+  %2 = dim %arg0, %c2 : tensor<?x?x?xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = linalg.generic
+    {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+                      affine_map<(d0, d1, d2) -> (d0, d2, d1)>,
+          affine_map<(d0, d1, d2) -> (d2, d1, d0)>],
+     iterator_types = ["parallel", "parallel", "parallel"]}
+    ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+    outs(%3 : tensor<?x?x?xf32>) {
+    ^bb0(%arg2 : f32, %arg3: f32, %arg4: f32):
+      %5 = addf %arg2, %arg3 : f32
+      linalg.yield %5 : f32
+    } -> tensor<?x?x?xf32>
+  return %4 : tensor<?x?x?xf32>
+}
+
+// CHECK-LABEL: func @generic_op_tensors
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+//       CHECK:   %[[INIT:.+]] = linalg.init_tensor
+//       CHECK:   %[[TD0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC0:.+]] = %[[INIT]]) -> (tensor<?x?x?xf32>) {
+//       CHECK:     %[[TD1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC1:.+]] = %[[TC0]]) -> (tensor<?x?x?xf32>) {
+//       CHECK:       %[[TD2:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC2:.+]] = %[[TC1]]) -> (tensor<?x?x?xf32>) {
+//       CHECK:       %[[STARG0:.+]] = subtensor %[[ARG0]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
+//       CHECK:       %[[STARG1:.+]] = subtensor %[[ARG1]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
+//       CHECK:       %[[STARG2:.+]] = subtensor %[[TC2]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
+//       CHECK:       %[[STRETURN:.+]] = linalg.generic
+//  CHECK-SAME:         ins(%[[STARG0]], %[[STARG1]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+//  CHECK-SAME:         outs(%[[STARG2]] : tensor<?x?x?xf32>)
+//       CHECK:       %[[TD:.+]] = subtensor_insert %[[STRETURN]] into %[[TC2]]
+//       CHECK:       scf.yield %[[TD]]
+//       CHECK:     }
+//       CHECK:     scf.yield %[[TD2]]
+//       CHECK:   }
+//       CHECK:   scf.yield %[[TD1]]
+//       CHECK: }
+//       CHECK: return %[[TD0]]
+
+// -----
+
+func @indexed_generic_op_tensors(
+  %arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %0 = dim %arg0, %c0 : tensor<?x?x?xf32>
+  %1 = dim %arg0, %c1 : tensor<?x?x?xf32>
+  %2 = dim %arg0, %c2 : tensor<?x?x?xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = linalg.indexed_generic
+    {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+                      affine_map<(d0, d1, d2) -> (d0, d2, d1)>,
+          affine_map<(d0, d1, d2) -> (d2, d1, d0)>],
+     iterator_types = ["parallel", "parallel", "parallel"]}
+    ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+    outs(%3 : tensor<?x?x?xf32>) {
+    ^bb0(%arg2 : index, %arg3 : index, %arg4 : index, %arg5 : f32, %arg6: f32, %arg7: f32):
+      %5 = addf %arg5, %arg6 : f32
+      linalg.yield %5 : f32
+    } -> tensor<?x?x?xf32>
+  return %4 : tensor<?x?x?xf32>
+}
+
+// CHECK-LABEL: func @indexed_generic_op_tensors
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
+//       CHECK:   %[[INIT:.+]] = linalg.init_tensor
+//       CHECK:   %[[TD0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC0:.+]] = %[[INIT]]) -> (tensor<?x?x?xf32>) {
+//       CHECK:     %[[TD1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC1:.+]] = %[[TC0]]) -> (tensor<?x?x?xf32>) {
+//       CHECK:       %[[TD2:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC2:.+]] = %[[TC1]]) -> (tensor<?x?x?xf32>) {
+//       CHECK:       %[[STARG0:.+]] = subtensor %[[ARG0]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
+//       CHECK:       %[[STARG1:.+]] = subtensor %[[ARG1]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
+//       CHECK:       %[[STARG2:.+]] = subtensor %[[TC2]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
+//       CHECK:       %[[STRETURN:.+]] = linalg.indexed_generic
+//  CHECK-SAME:         ins(%[[STARG0]], %[[STARG1]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+//  CHECK-SAME:         outs(%[[STARG2]] : tensor<?x?x?xf32>)
+//       CHECK:       %[[TD:.+]] = subtensor_insert %[[STRETURN]] into %[[TC2]]
+//       CHECK:       scf.yield %[[TD]]
+//       CHECK:     }
+//       CHECK:     scf.yield %[[TD2]]
+//       CHECK:   }
+//       CHECK:   scf.yield %[[TD1]]
+//       CHECK: }
+//       CHECK: return %[[TD0]]