[mlir][DO-NOT-REVIEW-IT-IS-A-DRAFT] updated getTiledImplementation

[pifon2a]

Summary

This is just to illustrate a comment in
https://discourse.llvm.org/t/rfc-changes-to-tilinginterface/66649

mlir/include/mlir/Interfaces/TilingInterface.h

 
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
 
+namespace mlir {
+
+struct TilingResult {
+  SmallVector<Operation *> generatedOperations;
+  SmallVector<Value> tiledResults;
+};
+
+struct SubsetSpecification {
+  struct Tile {
+    SmallVector<OpFoldResult> offsets;
+    SmallVector<OpFoldResult> sizes;
+    SmallVector<OpFoldResult> strides;
+  };
+  Tile tile;
+};
+
+using MaterializeSubsetFn = function_ref<Value(OpBuilder &, Location, Value,
+                                               const SubsetSpecification &)>;
+
+} // namespace mlir
+
 /// Include the ODS generated interface header files.
 #include "mlir/Interfaces/TilingInterface.h.inc"
 
 #endif // MLIR_INTERFACES_TILINGINTERFACE_H_

mlir/include/mlir/Interfaces/TilingInterface.td

let methods = [
             dimension had non-zero offset, it must be included in the offset
             provided here (as opposed to zero-based offset "relative" to the
             iteration space).
           - `sizes` provides the size of the tile.
 
           The method returns the operation that is the tiled
           implementation.
         }],
-        /*retType=*/"SmallVector<Operation *>",
+        /*retType=*/"TilingResult",
         /*methodName=*/"getTiledImplementation",
         /*args=*/(ins
             "OpBuilder &":$b,
-            "ArrayRef<OpFoldResult> ":$offsets,
-            "ArrayRef<OpFoldResult> ":$sizes),
+            "SubsetSpecification ":$subsetSpec,
+            "MaterializeSubsetFn ":$materializeFn),
         /*methodBody=*/"",
-        /*defaultImplementation=*/[{
-          return {};
-        }]
+        /*defaultImplementation=*/"return {};"
       >,
       InterfaceMethod<
         /*desc=*/[{
           Method to return the position of the result tile computed by the tiled operation.
 
           Specifies what tile of the result of the original tensor is computed
           by the tiled implementation. Expects the same `offsets` and `sizes` as
           used to obtain the tiled implementation of the operation.

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

static LogicalResult inlinePayload(OpBuilder &b, LinalgOp linalgOp,
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // External Model for implementing `TilingInterface` for `LinalgOp`s.
 //===----------------------------------------------------------------------===//
 
 namespace {
+
 /// External model implementation of TilingInterface for LinalgOps. An external
 /// model implementation is used for now till the use of `TilingInterface` is
 /// on-par with the current Linalg tiling + fusion patterns. Once it is
 /// maybe possible to move this into the op-definition (though there are
 /// advantages to leaving it as an external model)
 template <typename LinalgOpTy>
 struct LinalgOpTilingInterface
     : public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
return llvm::to_vector(
         llvm::map_range(map.getResults(), [&](AffineExpr loopExpr) {
           OpFoldResult ofr =
               makeComposedFoldedAffineApply(b, loc, loopExpr, allShapesSizes);
           return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)};
         }));
   }
 
   // Instantiate the tiled implementation of the operation.
-  SmallVector<Operation *>
-  getTiledImplementation(Operation *op, OpBuilder &b,
-                         ArrayRef<OpFoldResult> offsets,
-                         ArrayRef<OpFoldResult> sizes) const {
+  TilingResult getTiledImplementation(Operation *op, OpBuilder &b,
+                                      SubsetSpecification subsetSpec,
+                                      MaterializeSubsetFn materializeFn) 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->getOperands();
-    SmallVector<Value, 4> tiledOperands = makeTiledShapes(
-        b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
+
+    SmallVector<Optional<SliceParameters>> allSliceParameter =
+        computeAllSliceParameters(b, loc, linalgOp, valuesToTile,
+                                  subsetSpec.tile.offsets,
+                                  subsetSpec.tile.strides, {}, true);
+
+    SmallVector<Value> tiledOperands;
+    for (auto item : llvm::zip(valuesToTile, allSliceParameter)) {
+      Value valueToTile = std::get<0>(item);
+      Optional<SliceParameters> sliceParams = std::get<1>(item);
+      tiledOperands.push_back(
+          sliceParams.has_value()
+              ? materializeFn(b, loc, valueToTile,
+                              {{sliceParams->offsets, sliceParams->sizes,
+                                sliceParams->strides}})
+              : valueToTile);
+    }
 
     SmallVector<Type> resultTensorTypes =
         getTensorOutputTypes(linalgOp, tiledOperands);
 
     Operation *tiledOp =
         linalgOp.clone(b, loc, resultTensorTypes, tiledOperands);
-    offsetIndices(b, cast<LinalgOp>(tiledOp), offsets);
+    offsetIndices(b, cast<LinalgOp>(tiledOp), subsetSpec.tile.offsets);
 
-    return {tiledOp};
+    return TilingResult{{tiledOp}, tiledOp->getResults()};
   }
 
   // Return the details of the output tile generated by the tiled
   // implementation.
   LogicalResult
   getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                         ArrayRef<OpFoldResult> offsets,
                         ArrayRef<OpFoldResult> sizes,

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

updateDestinationOperandsForTiledOp(OpBuilder &builder,
   for (const auto &destValue : llvm::enumerate(tiledOpDestinationValues)) {
     auto sliceOp = destValue.value().getDefiningOp<tensor::ExtractSliceOp>();
     if (!sliceOp)
       continue;
     sliceOp.setOperand(0, bbArgsList[destValue.index()]);
   }
 }
 
+Value MaterializeSlice(OpBuilder &builder, Location loc, Value valueToTile,
+                       const SubsetSpecification &subsetSpec) {
+  return builder.create<tensor::ExtractSliceOp>(
+      loc, valueToTile, subsetSpec.tile.offsets, subsetSpec.tile.sizes,
+      subsetSpec.tile.strides);
+}
+
 /// Implementation of tiling transformation of `op` that implements the
 /// `TilingInterface` using `scf.for` to iterate over the tiles.
 FailureOr<scf::SCFTilingResult>
 mlir::scf::tileUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
                              const scf::SCFTilingOptions &options) {
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPointAfter(op);
 
if (!tilingResult.loops.empty()) {
       llvm::dbgs() << "\n";
     }
   });
 
   // 4. Generate the tiled implementation within the inner most loop.
   if (!tilingResult.loops.empty())
     rewriter.setInsertionPoint(
         tilingResult.loops.back().getBody()->getTerminator());
-  SmallVector<Operation *> tiledImplementation =
-      op.getTiledImplementation(rewriter, offsets, sizes);
-  tilingResult.tiledOps.append(tiledImplementation);
+  SubsetSpecification subsetSpec;
+  subsetSpec.tile = {
+      offsets, sizes,
+      SmallVector<OpFoldResult>(offsets.size(), rewriter.getIndexAttr(1))};
+
+  TilingResult tiledImplementationResult =
+      op.getTiledImplementation(rewriter, subsetSpec, MaterializeSlice);
+  tilingResult.tiledOps.append(tiledImplementationResult.generatedOperations);
   if (op->getNumResults() == 0) {
     // nothing more to do.
     return tilingResult;
   }
 
   // If loops are empty, the tiled op is used as the replacement for the untiled
   // op.
   if (tilingResult.loops.empty()) {
-    tilingResult.replacements = llvm::to_vector(
-        llvm::map_range(tiledImplementation[0]->getResults(),
-                        [](OpResult result) -> Value { return result; }));
+    tilingResult.replacements = tiledImplementationResult.tiledResults;
     return tilingResult;
   }
 
   // 5. Yield all the results of the tiled operation. The surrounding loop
   //    nest is modified to insert a destructive update pattern to yield
   //    from the loop nest values to replace the untiled op with.
   int64_t numResults = op->getNumResults();
   SmallVector<SmallVector<OpFoldResult>> resultOffsetsList(numResults),