[mlir][linalg] Remove IndexedGenericOp support from LinalgInterchangePattern...

[gysit]

Summary

after introducing the IndexedGenericOp to GenericOp canonicalization (https://reviews.llvm.org/D101612).

[gysit]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

638–641

Replacing the LinalgInterchangePattern<GenericOp> by the GenericOpInterchangePattern is maybe more change than need and affects downstream projects. Instead we could stick LinalgInterchangePattern and support only GenericOps...

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// the index accesses of `op`. This is an in-place transformation controlled by
 /// `interchangeVector`. An empty vector is interpreted as the identity
 /// permutation and the transformation returns early.
 ///
 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed with
 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
 /// integers, in the range 0..`op.rank` without duplications
 /// (i.e. `[1,1,2]` is an invalid permutation).
-void interchange(PatternRewriter &rewriter, LinalgOp op,
+void interchangeGenericOp(PatternRewriter &rewriter, GenericOp genericOp,
                           ArrayRef<unsigned> interchangeVector);
 
 /// Callback function type used to perform the allocation for the promoted
 /// `subView`. In `boundingSubViewsize` a best attempt is made to find the
 /// smallest constant value for the size of the buffer needed for each
 /// dimension. If that is not possible, contains the dynamic size of the
 /// subview. The call back should return the buffer to use.
 using AllocBufferCallbackFn =
     std::function<Optional<Value>(OpBuilder &b, memref::SubViewOp subView,

 /// Emits a loop nest of `affine.for` with the proper body for `linalgOp`.
 LogicalResult linalgOpToAffineLoops(PatternRewriter &rewriter,
                                     LinalgOp linalgOp);
 
 //===----------------------------------------------------------------------===//
 // Preconditions that ensure the corresponding transformation succeeds and can
 // be applied as a rewrite pattern.
 //===----------------------------------------------------------------------===//
-/// Emits a `generic` or `indexed_generic` operation with the `indexing_maps`
-/// and `iterator_types` permutated according to `permutation`.
+/// Emits a `generic` operation with the `indexing_maps` and `iterator_types`
+/// permutated according to `permutation`.
 LogicalResult
-interchangeGenericLinalgOpPrecondition(Operation *op,
-                                       ArrayRef<unsigned> interchangeVector);
+interchangeGenericOpPrecondition(GenericOp genericOp,
 
 /// Promote std.subviews feeding linalg operations.
 LogicalResult promoteSubviewsPrecondition(Operation *op,
                                           LinalgPromotionOptions options);
 
 /// Rewrite a linalg.generic into a suitable vector.contraction op.
 LogicalResult vectorizeLinalgOpPrecondition(Operation *op);
 
LinalgTileAndFusePattern(
           LinalgTransformationFilter(),
       PatternBenefit benefit = 1)
       : LinalgBaseTileAndFusePattern(
             OpTy::getOperationName(), context, dependenceGraph, tilingOptions,
             fusionOptions, filter, fusedOpMarker, originalOpMarker, benefit) {}
 };
 
 ///
-/// Linalg interchange patterns.
+/// Linalg generic interchage pattern.
 ///
 /// Apply the `interchange` transformation as a pattern.
 /// `filter` controls LinalgTransformMarker matching and update when specified.
 /// See `interchange` for more details.
-struct LinalgBaseInterchangePattern : public RewritePattern {
-  LinalgBaseInterchangePattern(
-      StringRef opName, MLIRContext *context,
-      ArrayRef<unsigned> interchangeVector,
+struct GenericOpInterchangePattern : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+  GenericOpInterchangePattern(
+      MLIRContext *context, ArrayRef<unsigned> interchangeVector,
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       PatternBenefit benefit = 1);
-  LogicalResult matchAndRewrite(Operation *op,
+  LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override;
 
 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgTransformationFilter filter;
   /// The interchange vector to reorder the iterators and indexing_maps dims.
   SmallVector<unsigned, 8> interchangeVector;
 };
 
-template <typename OpTy>
-struct LinalgInterchangePattern : public LinalgBaseInterchangePattern {
-  LinalgInterchangePattern(
-      MLIRContext *context, ArrayRef<unsigned> interchangeVector,
-      LinalgTransformationFilter filter = LinalgTransformationFilter(),
-      PatternBenefit benefit = 1)
-      : LinalgBaseInterchangePattern(OpTy::getOperationName(), context,
-                                     interchangeVector, filter, benefit) {}
-};
-
 ///
 /// Linalg promotion patterns.
 ///
 /// Apply the `promoteSubViews` transformation as a pattern.
 /// `filter` controls LinalgTransformMarker matching and update when specified.
 /// See `promoteSubViews` for more details.
 struct LinalgBasePromotionPattern : public RewritePattern {
   LinalgBasePromotionPattern(

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

 #include "llvm/Support/raw_ostream.h"
 #include <type_traits>
 
 #define DEBUG_TYPE "linalg-interchange"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
-LogicalResult mlir::linalg::interchangeGenericLinalgOpPrecondition(
-    Operation *op, ArrayRef<unsigned> interchangeVector) {
-  // Transformation applies to generic ops only.
-  if (!isa<GenericOp, IndexedGenericOp>(op))
-    return failure();
-  LinalgOp linalgOp = cast<LinalgOp>(op);
+LogicalResult mlir::linalg::interchangeGenericOpPrecondition(
+    GenericOp genericOp, ArrayRef<unsigned> interchangeVector) {
   // Interchange vector must be non-empty and match the number of loops.
   if (interchangeVector.empty() ||
-      linalgOp.getNumLoops() != interchangeVector.size())
+      genericOp.getNumLoops() != interchangeVector.size())
     return failure();
   // Permutation map must be invertible.
-  if (!inversePermutation(
-          AffineMap::getPermutationMap(interchangeVector, op->getContext())))
+  if (!inversePermutation(AffineMap::getPermutationMap(interchangeVector,
+                                                       genericOp.getContext())))
     return failure();
   return success();
 }
 
-void mlir::linalg::interchange(PatternRewriter &rewriter, LinalgOp op,
+void mlir::linalg::interchangeGenericOp(PatternRewriter &rewriter,
+                                        GenericOp genericOp,
                                         ArrayRef<unsigned> interchangeVector) {
   // 1. Compute the inverse permutation map.
-  MLIRContext *context = op.getContext();
+  MLIRContext *context = genericOp.getContext();
   AffineMap permutationMap = inversePermutation(
       AffineMap::getPermutationMap(interchangeVector, context));
   assert(permutationMap && "expected permutation to be invertible");
-  assert(interchangeVector.size() == op.getNumLoops() &&
+  assert(interchangeVector.size() == genericOp.getNumLoops() &&
          "expected interchange vector to have entry for every loop");
 
   // 2. Compute the interchanged indexing maps.
   SmallVector<Attribute, 4> newIndexingMaps;
-  ArrayRef<Attribute> indexingMaps = op.indexing_maps().getValue();
-  for (unsigned i = 0, e = op.getNumShapedOperands(); i != e; ++i) {
+  ArrayRef<Attribute> indexingMaps = genericOp.indexing_maps().getValue();
+  for (unsigned i = 0, e = genericOp.getNumShapedOperands(); i != e; ++i) {
     AffineMap m = indexingMaps[i].cast<AffineMapAttr>().getValue();
     if (!permutationMap.isEmpty())
       m = m.compose(permutationMap);
     newIndexingMaps.push_back(AffineMapAttr::get(m));
   }
-  op->setAttr(getIndexingMapsAttrName(),
+  genericOp->setAttr(getIndexingMapsAttrName(),
                      ArrayAttr::get(context, newIndexingMaps));
 
   // 3. Compute the interchanged iterator types.
-  ArrayRef<Attribute> itTypes = op.iterator_types().getValue();
+  ArrayRef<Attribute> itTypes = genericOp.iterator_types().getValue();
   SmallVector<Attribute, 4> itTypesVector;
   llvm::append_range(itTypesVector, itTypes);
   applyPermutationToVector(itTypesVector, interchangeVector);
-  op->setAttr(getIteratorTypesAttrName(),
+  genericOp->setAttr(getIteratorTypesAttrName(),
                      ArrayAttr::get(context, itTypesVector));
 
   // 4. Transform the index operations by applying the permutation map.
-  if (op.hasIndexSemantics()) {
+  if (genericOp.hasIndexSemantics()) {
     // TODO: Remove the assertion and add a getBody() method to LinalgOp
     // interface once every LinalgOp has a body.
-    assert(op->getNumRegions() == 1 &&
-           op->getRegion(0).getBlocks().size() == 1 &&
+    assert(genericOp->getNumRegions() == 1 &&
+           genericOp->getRegion(0).getBlocks().size() == 1 &&
            "expected generic operation to have one block.");
-    Block &block = op->getRegion(0).front();
+    Block &block = genericOp->getRegion(0).front();
     OpBuilder::InsertionGuard guard(rewriter);
     for (IndexOp indexOp :
          llvm::make_early_inc_range(block.getOps<IndexOp>())) {
       rewriter.setInsertionPoint(indexOp);
       SmallVector<Value> allIndices;
-      allIndices.reserve(op.getNumLoops());
-      llvm::transform(llvm::seq<uint64_t>(0, op.getNumLoops()),
+      allIndices.reserve(genericOp.getNumLoops());
+      llvm::transform(llvm::seq<uint64_t>(0, genericOp.getNumLoops()),
                       std::back_inserter(allIndices), [&](uint64_t dim) {
                         return rewriter.create<IndexOp>(indexOp->getLoc(), dim);
                       });
       rewriter.replaceOpWithNewOp<AffineApplyOp>(
           indexOp, permutationMap.getSubMap(indexOp.dim()), allIndices);
     }
   }
 }

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

originalOpMarker.replaceLinalgTransformationFilter(
         rewriter, origProducerOp.getOperation());
   }
   rewriter.updateRootInPlace(op, [&]() {
     originalOpMarker.replaceLinalgTransformationFilter(rewriter, op);
   });
   return success();
 }
 
-/// Linalg base interchange pattern.
-mlir::linalg::LinalgBaseInterchangePattern::LinalgBaseInterchangePattern(
-    StringRef opName, MLIRContext *context,
-    ArrayRef<unsigned> interchangeVector, LinalgTransformationFilter filter,
-    PatternBenefit benefit)
-    : RewritePattern(opName, benefit, context, {}), filter(filter),
+/// Linalg generic interchange pattern.
+mlir::linalg::GenericOpInterchangePattern::GenericOpInterchangePattern(
+    MLIRContext *context, ArrayRef<unsigned> interchangeVector,
+    LinalgTransformationFilter filter, PatternBenefit benefit)
+    : OpRewritePattern(context, benefit), filter(filter),
       interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
 
-LogicalResult mlir::linalg::LinalgBaseInterchangePattern::matchAndRewrite(
-    Operation *op, PatternRewriter &rewriter) const {
-  LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-  if (!linalgOp)
-    return failure();
-  if (failed(filter.checkAndNotify(rewriter, linalgOp)))
+LogicalResult mlir::linalg::GenericOpInterchangePattern::matchAndRewrite(
+    GenericOp genericOp, PatternRewriter &rewriter) const {
+  if (failed(filter.checkAndNotify(rewriter, genericOp)))
     return failure();
-  if (failed(interchangeGenericLinalgOpPrecondition(op, interchangeVector)))
+  if (failed(interchangeGenericOpPrecondition(genericOp, interchangeVector)))
     return failure();
 
   // TODO: figure out how this interplays with named ops. In particular this
   // should break the named op property.
-  rewriter.updateRootInPlace(op, [&]() {
-    interchange(rewriter, linalgOp, interchangeVector);
+  rewriter.updateRootInPlace(genericOp, [&]() {
+    interchangeGenericOp(rewriter, genericOp, interchangeVector);
     // New filter if specified.
-    filter.replaceLinalgTransformationFilter(rewriter, op);
+    filter.replaceLinalgTransformationFilter(rewriter, genericOp);
   });
   return success();
 }
 
 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
     StringRef opName, MLIRContext *context, LinalgPromotionOptions options,
     LinalgTransformationFilter filter, PatternBenefit benefit)
     : RewritePattern(opName, benefit, context, {}), filter(filter),

mlir/test/Dialect/Linalg/transform-patterns.mlir

 // CHECK:        linalg.generic {
 // CHECK-SAME:   indexing_maps = [#[[$kn]], #[[$nm]], #[[$km]]],
 // CHECK-SAME:   iterator_types = ["parallel", "reduction", "parallel"],
 // CHECK-SAME:   library_call = "linalg_matmul"}
 // CHECK:          memref<?x?xf32, #[[$STRIDED_2D_u_1]]>,
 // CHECK-SAME:     memref<?x?xf32, #[[$STRIDED_2D_u_1]]>
 // CHECK-SAME:     memref<?x?xf32, #[[$STRIDED_2D_u_1]]>
 
-#indexed_matmul_trait = {
-  args_in = 2,
-  args_out = 1,
-  indexing_maps = #matmul_accesses,
-  library_call = "linalg_matmul_indexed",
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-func @permute_generic_indexed(
-    %A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
-    %B: memref<?x?xf32, offset: ?, strides: [?, 1]>,
-    %C: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
-  linalg.indexed_generic #indexed_matmul_trait
-    ins(%A, %B : memref<?x?xf32, offset: ?, strides: [?, 1]>,
-                 memref<?x?xf32, offset: ?, strides: [?, 1]>)
-   outs(%C : memref<?x?xf32, offset: ?, strides: [?, 1]>) {
-    ^bb(%i: index, %j: index, %k: index, %a: f32, %b: f32, %c: f32):
-      %d = mulf %a, %b: f32
-      %e = addf %c, %d: f32
-      linalg.yield %e: f32
-  }
-  return
-}
-// CHECK-LABEL:  func @permute_generic_indexed
-// CHECK:        linalg.indexed_generic {
-// CHECK-SAME:     indexing_maps = [#[[$kn]], #[[$nm]], #[[$km]]],
-// CHECK-SAME:     iterator_types = ["parallel", "reduction", "parallel"],
-// CHECK-SAME:     library_call = "linalg_matmul_indexed"}
-// CHECK:            memref<?x?xf32, #[[$STRIDED_2D_u_1]]>,
-// CHECK-SAME:       memref<?x?xf32, #[[$STRIDED_2D_u_1]]>
-// CHECK-SAME:       memref<?x?xf32, #[[$STRIDED_2D_u_1]]>
-
 func @matvec_perm(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
              %x: memref<?xf32, offset: ?, strides: [1]>,
              %y: memref<?xf32, offset: ?, strides: [1]>) {
   linalg.matvec {__internal_linalg_transform__ = "__with_perm__"}
     ins(%A, %x: memref<?x?xf32, offset: ?, strides: [?, 1]>,
                 memref<?xf32, offset: ?, strides: [1]>)
    outs(%y: memref<?xf32, offset: ?, strides: [1]>)
   return

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

static void applyPatterns(FuncOp funcOp) {
   //===--------------------------------------------------------------------===//
   // Linalg to vector contraction patterns.
   //===--------------------------------------------------------------------===//
   patterns.add<LinalgVectorizationPattern>(
       ctx, LinalgTransformationFilter(Identifier::get("VECTORIZE", ctx))
                .addOpFilter<MatmulOp, FillOp, CopyOp, GenericOp>());
 
   //===--------------------------------------------------------------------===//
-  // Linalg generic permutation patterns.
+  // Linalg generic interchange pattern.
   //===--------------------------------------------------------------------===//
-  patterns.add<LinalgInterchangePattern<GenericOp>>(
-      ctx,
-      /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
-      LinalgTransformationFilter(ArrayRef<Identifier>{},
-                                 Identifier::get("PERMUTED", ctx)));
-  patterns.add<LinalgInterchangePattern<IndexedGenericOp>>(
+  patterns.add<GenericOpInterchangePattern>(
       ctx,
       /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
       LinalgTransformationFilter(ArrayRef<Identifier>{},
                                  Identifier::get("PERMUTED", ctx)));
 
   //===--------------------------------------------------------------------===//
   // Linalg subview operands promotion.
   //===--------------------------------------------------------------------===//
tilingPattern.add<linalg::LinalgTilingPattern<linalg::MatmulI8I8I32Op>>(
           Identifier::get("tile-and-pad", context)));
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
 }
 
 static void applyInterchangePattern(FuncOp funcOp,
                                     ArrayRef<unsigned> interchangeVector) {
   MLIRContext *context = funcOp.getContext();
   RewritePatternSet interchangePattern(context);
-  interchangePattern.add<LinalgInterchangePattern<GenericOp>>(
+  interchangePattern.add<GenericOpInterchangePattern>(
       context, interchangeVector,
       LinalgTransformationFilter(ArrayRef<Identifier>{},
                                  Identifier::get("interchange", context)));
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(interchangePattern));
 }
 
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnFunction() {