[mlir][linalg] adding pass to run the interchange pattern.

[gysit]

Summary

Instead of interchanging loops during the loop lowering this pass performs the interchange by permuting the indexing maps. It also updates the iterator types and the index accesses in the body of the operation.

[gysit]

The pass uses the existing LinalgInterchangePattern extended with the functionality to update the index operations in the body of the interchanged generic op. I relaxed the constraint that the pass can only run ops with buffer semantics but kept the limitation on generic operations. Should we relax these constraints to support all linalgOps?

[nicolasvasilache]

I'd think there is very little chance that this pass is going to be useful in real compilation flows (granularity of control, need for heurisitcs, phase orderings etc etc).
I'd rather add a new case + attribute in the test pass that tests individual patterns and call it a day.

[gysit]

Add the pattern as option to the test pass instead of having a separate interchange pass.

[gysit]

Update to current head.

[nicolasvasilache]

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

90

op.getBlock should do it, the verifier ensures that the region is there.
Does the assertion trigger for you?
I thought all named ops now have a region but I may be wrong ..

93

I'd lift that one level or even to the top of the function.

[gysit]

Address comments.

[gysit]

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

90

Added a TODO to fix this once all LinalgOps have a body

[nicolasvasilache]

Nice!

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

207–208

How about just taking an OpBuilder here ?
That would be more generally applicable.

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

68

These days, we should be able to simplify this with something like

SmallVector<AffineMap> newIndexingMaps;
for (auto m : op.indexing_maps().getAsValueRange<AffineMap>()) {
  if (!permutationMap.isEmpty())
      m = m.compose(permutationMap);
    newIndexingMaps.push_back(m);
}
 op->setAttr(getIndexingMapsAttrName(), rewriter.getAffineMapArrayAttr(newIndexingMaps));

90

yes sorry this is a bad suggestion, getBlock() gives you the parent block, please disregard.

mlir/include/mlir/Dialect/Linalg/Passes.h

 #define MLIR_DIALECT_LINALG_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 std::unique_ptr<OperationPass<FuncOp>> createConvertElementwiseToLinalgPass();
 
 std::unique_ptr<OperationPass<FuncOp>> createLinalgFoldUnitExtentDimsPass();
+std::unique_ptr<OperationPass<FuncOp>> createLinalgInterchangePass();
 
 std::unique_ptr<Pass> createLinalgFusionOfTensorOpsPass();
 std::unique_ptr<Pass> createFoldReshapeOpsByLinearizationPass();
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>

mlir/include/mlir/Dialect/Linalg/Passes.td

Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
            "Only folds the one-trip loops from Linalg ops on tensors "
            "(for testing purposes only)">
   ];
   let dependentDialects = [
     "linalg::LinalgDialect", "AffineDialect", "memref::MemRefDialect"
   ];
 }
 
+def LinalgInterchange : FunctionPass<"linalg-interchange"> {
+  let summary = "Interchange the iteration domain basis of linalg operations";
+  let constructor = "mlir::createLinalgInterchangePass()";
+  let options = [
+    ListOption<"interchangeVector", "interchange-vector", "unsigned",
+               "Permute the dimensions of the iteration domain following the "
+               "the given interchange vector",
+               "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
+  ];
+  let dependentDialects = [
+    "linalg::LinalgDialect",
+    "AffineDialect"
+  ];
+}
+
 def LinalgFusionOfTensorOps : Pass<"linalg-fusion-for-tensor-ops"> {
   let summary = "Fuse operations on RankedTensorType in linalg dialect";
   let constructor = "mlir::createLinalgFusionOfTensorOpsPass()";
   let options = [
     Option<"allowFoldingUnitDimReshapes", "allow-folding-unit-dim-reshapes",
            "bool", /*default=*/"false",
            "Allow fusing linalg.tensor_reshape ops that performs unit "
            "dimension collapsing">

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

struct TiledAndFusedLinalgOps {
   SmallVector<Operation *, 4> fusedLoops;
 };
 Optional<TiledAndFusedLinalgOps>
 tileAndFuseLinalgOps(OpBuilder &builder, ArrayRef<LinalgOp> ops,
                      const LinalgDependenceGraph &dependenceGraph,
                      const LinalgTilingOptions &tilingOptions);
 
 /// Interchanges the `iterator_types` and `iterator_maps` dimensions 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.
+/// Additionally, the updates the dimensions of all index op accesses in the
+/// body of transformed operation. 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).
-LinalgOp interchange(LinalgOp op, ArrayRef<unsigned> interchangeVector);
+void interchange(PatternRewriter &rewriter, LinalgOp op,
+                 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,

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

 //===- Interchange.cpp - Linalg interchange transformation ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the linalg interchange transformation.
 //
 //===----------------------------------------------------------------------===//
 
+#include "PassDetail.h"
 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/Debug.h"
 #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) {
-  if (interchangeVector.empty())
-    return failure();
   // Transformation applies to generic ops only.
   if (!isa<GenericOp, IndexedGenericOp>(op))
     return failure();
-  LinalgOp linOp = cast<LinalgOp>(op);
-  // Transformation applies to buffers only.
-  if (!linOp.hasBufferSemantics())
-    return failure();
-  // Permutation must be applicable.
-  if (linOp.getIndexingMap(0).getNumInputs() != interchangeVector.size())
+  LinalgOp linalgOp = cast<LinalgOp>(op);
+  // Interchange vector must be non-empty and match the number of loops.
+  if (interchangeVector.empty() ||
+      linalgOp.getNumLoops() != interchangeVector.size())
     return failure();
   // Permutation map must be invertible.
   if (!inversePermutation(
           AffineMap::getPermutationMap(interchangeVector, op->getContext())))
     return failure();
   return success();
 }
 
-LinalgOp mlir::linalg::interchange(LinalgOp op,
-                                   ArrayRef<unsigned> interchangeVector) {
-  if (interchangeVector.empty())
-    return op;
-
+void mlir::linalg::interchange(PatternRewriter &rewriter, LinalgOp op,
+  // 1. Compute the inverse permutation map.
   MLIRContext *context = op.getContext();
-  auto permutationMap = inversePermutation(
+  AffineMap permutationMap = inversePermutation(
       AffineMap::getPermutationMap(interchangeVector, context));
   assert(permutationMap && "expected permutation to be invertible");
+  assert(interchangeVector.size() == op.getNumLoops() &&
+         "expected interchange vector to have entry for every loop");
+
+  // 2. Compute the interchanged indexing maps.
   SmallVector<Attribute, 4> newIndexingMaps;
-  auto indexingMaps = op.indexing_maps().getValue();
+  ArrayRef<Attribute> indexingMaps = op.indexing_maps().getValue();
   for (unsigned i = 0, e = op.getNumShapedOperands(); i != e; ++i) {
     AffineMap m = indexingMaps[i].cast<AffineMapAttr>().getValue();
     if (!permutationMap.isEmpty())
       m = m.compose(permutationMap);
     newIndexingMaps.push_back(AffineMapAttr::get(m));
   }
-  auto itTypes = op.iterator_types().getValue();
-  SmallVector<Attribute, 4> itTypesVector;
-  for (unsigned i = 0, e = itTypes.size(); i != e; ++i)
-    itTypesVector.push_back(itTypes[i]);
-  applyPermutationToVector(itTypesVector, interchangeVector);
-
   op->setAttr(getIndexingMapsAttrName(),
               ArrayAttr::get(context, newIndexingMaps));
+
+  // 3. Compute the interchanged iterator types.
+  ArrayRef<Attribute> itTypes = op.iterator_types().getValue();
+  SmallVector<Attribute, 4> itTypesVector;
+  llvm::append_range(itTypesVector, itTypes);
+  applyPermutationToVector(itTypesVector, interchangeVector);
   op->setAttr(getIteratorTypesAttrName(),
               ArrayAttr::get(context, itTypesVector));
 
-  return op;
+  // 4. Transform the index operations by applying the permutation map.
+  if (op.hasIndexSemantics()) {
+    assert(op->getNumRegions() == 1 &&
+           op->getRegion(0).getBlocks().size() == 1 &&
+           "expected generic operation to have one block.");
+    Block &block = op->getRegion(0).front();
+    for (IndexOp indexOp :
+         llvm::make_early_inc_range(block.getOps<IndexOp>())) {
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPointAfter(indexOp);
+      SmallVector<Value> allIndices;
+      allIndices.reserve(op.getNumLoops());
+      llvm::transform(llvm::seq<int64_t>(0, op.getNumLoops()),
+                      std::back_inserter(allIndices), [&](int64_t dim) {
+                        return rewriter.create<IndexOp>(indexOp.getLoc(), dim);
+                      });
+      rewriter.replaceOpWithNewOp<AffineApplyOp>(
+          indexOp, permutationMap.getSubMap(indexOp.dim()), allIndices);
+    }
+  }
+}
+
+namespace {
+
+/// Pass that interchanges the iteration domain basis of linalg operations.
+struct LinalgInterchangePass
+    : public LinalgInterchangeBase<LinalgInterchangePass> {
+  void runOnFunction() override {
+    FuncOp funcOp = getFunction();
+    MLIRContext *context = funcOp.getContext();
+
+    // Apply the interchange pattern to all generic operations.
+    RewritePatternSet patterns(context);
+    patterns.add<LinalgInterchangePattern<GenericOp>>(
+        context, interchangeVector,
+        LinalgTransformationFilter(ArrayRef<Identifier>{},
+                                   Identifier::get("interchange", context)));
+    (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+
+    // Cleanup all transformation markers.
+    funcOp.walk([](LinalgOp op) {
+      op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
+    });
+  }
+};
+} // namespace
+
+std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgInterchangePass() {
+  return std::make_unique<LinalgInterchangePass>();
 }

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

mlir::linalg::LinalgBaseInterchangePattern::LinalgBaseInterchangePattern(
     ArrayRef<unsigned> interchangeVector, LinalgTransformationFilter filter,
     PatternBenefit benefit)
     : RewritePattern(opName, benefit, context, {}), filter(filter),
       interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
 
 LogicalResult mlir::linalg::LinalgBaseInterchangePattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-  // TODO: remove hasIndexSemantics check once index ops are supported.
-  if (!linalgOp || linalgOp.hasIndexSemantics())
+  if (!linalgOp)
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();
   if (failed(interchangeGenericLinalgOpPrecondition(op, 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(linalgOp, interchangeVector);
+    interchange(rewriter, linalgOp, interchangeVector);
     // New filter if specified.
     filter.replaceLinalgTransformationFilter(rewriter, op);
   });
   return success();
 }
 
 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
     StringRef opName, MLIRContext *context, LinalgPromotionOptions options,

mlir/test/Dialect/Linalg/interchange.mlir

+// RUN: mlir-opt %s -linalg-interchange="interchange-vector=4,0,3,1,2" | FileCheck %s
+// RUN: mlir-opt %s -linalg-interchange="interchange-vector=4,0,3,1,2" -linalg-interchange="interchange-vector=1,3,4,2,0"  | FileCheck --check-prefix=CANCEL-OUT %s
+
+#map0 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
+#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3)>
+
+func @interchange_generic_op(%arg0 : memref<1x2x3x4x5xindex>, %arg1 : memref<1x2x4xindex>) {
+  linalg.generic {
+    indexing_maps = [#map0, #map1],
+    iterator_types = ["parallel", "parallel", "reduction", "parallel", "reduction"]}
+  ins(%arg0 : memref<1x2x3x4x5xindex>)
+  outs(%arg1 : memref<1x2x4xindex>) {
+      ^bb0(%arg2 : index, %arg3 : index) :
+        %0 = linalg.index 0 : index
+        %1 = linalg.index 1 : index
+        %2 = linalg.index 4 : index
+        %3 = subi %0, %1 : index
+        %4 = addi %3, %2 : index
+        %5 = addi %4, %arg2 : index
+        linalg.yield %5 : index
+      }
+  return
+}
+
+//    CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3, d4, d2, d0)>
+//    CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3, d2)>
+//        CHECK: func @interchange_generic_op
+//        CHECK:   linalg.generic
+//   CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
+//   CHECK-SAME:     iterator_types = ["reduction", "parallel", "parallel", "parallel", "reduction"]
+//    CHECK-DAG:     %[[IDX0:.+]] = linalg.index 1 : index
+//    CHECK-DAG:     %[[IDX1:.+]] = linalg.index 3 : index
+//    CHECK-DAG:     %[[IDX4:.+]] = linalg.index 0 : index
+//        CHECK:     %[[T0:.+]] = subi %[[IDX0]], %[[IDX1]] : index
+//        CHECK:     %[[T1:.+]] = addi %[[T0]], %[[IDX4]] : index
+//        CHECK:     %[[T2:.+]] = addi %[[T1]], %{{.*}} : index
+
+//  CANCEL-OUT-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
+//  CANCEL-OUT-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3)>
+//      CANCEL-OUT: func @interchange_generic_op
+//      CANCEL-OUT:   linalg.generic
+// CANCEL-OUT-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
+// CANCEL-OUT-SAME:     iterator_types = ["parallel", "parallel", "reduction", "parallel", "reduction"]
+//  CANCEL-OUT-DAG:     %[[IDX0:.+]] = linalg.index 0 : index
+//  CANCEL-OUT-DAG:     %[[IDX1:.+]] = linalg.index 1 : index
+//  CANCEL-OUT-DAG:     %[[IDX4:.+]] = linalg.index 4 : index
+//      CANCEL-OUT:     %[[T0:.+]] = subi %[[IDX0]], %[[IDX1]] : index
+//      CANCEL-OUT:     %[[T1:.+]] = addi %[[T0]], %[[IDX4]] : index
+//      CANCEL-OUT:     %[[T2:.+]] = addi %[[T1]], %{{.*}} : index
+
+