[mlir][Vector] Modernize default lowering of vector transpose

[dcaballe]

Summary

This patch removes an old recursive implementation to lower vector.transpose to extract/insert operations
and replaces it with a iterative approach that leverages newer linearization/delinearization utilities.
The patch should be NFC except by the order in which the extract/insert ops are generated.

[nicolasvasilache]

mlir/lib/Dialect/Utils/IndexingUtils.cpp

33 ↗

(On Diff #414176)

actually, now that I see this, I remember we have:

void applyPermutationToVector(SmallVector<T, N> &inVec,
                              ArrayRef<int64_t> permutation)

in LinalgUtils.h; mind refactoring a bit more so that one uses the other / we end with only 1 version ?

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

421

can we use/add a builder that would take ArrayRef and avoid the need for users to materialize the Attr ?

[dcaballe]

Addressed feedback.

[dcaballe]

I'll commit it tomorrow if no more comments. Thanks!

mlir/lib/Dialect/Utils/IndexingUtils.cpp

33 ↗

(On Diff #414176)

I moved applyPermutationToVector to this file, since it's more generic and can be reused across dialects.

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

421

Much better, thanks!

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 //===----------------------------------------------------------------------===//
 // General utilities
 //===----------------------------------------------------------------------===//
 
 /// Check if `permutation` is a permutation of the range
 /// `[0, permutation.size())`.
 bool isPermutation(ArrayRef<int64_t> permutation);
 
-/// Apply the permutation defined by `permutation` to `inVec`.
-/// Element `i` in `inVec` is mapped to location `j = permutation[i]`.
-/// E.g.: for an input vector `inVec = ['a', 'b', 'c']` and a permutation vector
-/// `permutation = [2, 0, 1]`, this function leaves `inVec = ['c', 'a', 'b']`.
-template <typename T, unsigned N>
-void applyPermutationToVector(SmallVector<T, N> &inVec,
-                              ArrayRef<int64_t> permutation) {
-  SmallVector<T, N> auxVec(inVec.size());
-  for (const auto &en : enumerate(permutation))
-    auxVec[en.index()] = inVec[en.value()];
-  inVec = auxVec;
-}
-
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
 
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.
 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
 

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

 int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis);
 
 /// Given the strides together with a linear index in the dimension
 /// space, returns the vector-space offsets in each dimension for a
 /// de-linearized index.
 SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> strides,
                                     int64_t linearIndex);
 
+/// Apply the permutation defined by `permutation` to `inVec`.
+/// Element `i` in `inVec` is mapped to location `j = permutation[i]`.
+/// E.g.: for an input vector `inVec = ['a', 'b', 'c']` and a permutation vector
+/// `permutation = [2, 0, 1]`, this function leaves `inVec = ['c', 'a', 'b']`.
+template <typename T, unsigned N>
+void applyPermutationToVector(SmallVector<T, N> &inVec,
+                              ArrayRef<int64_t> permutation) {
+  SmallVector<T, N> auxVec(inVec.size());
+  for (const auto &en : enumerate(permutation))
+    auxVec[en.index()] = inVec[en.value()];
+  inVec = auxVec;
+}
+
 /// Helper that returns a subset of `arrayAttr` as a vector of int64_t.
 SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr,
                                        unsigned dropFront = 0,
                                        unsigned dropBack = 0);
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_INDEXINGUTILS_H

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

 #include "PassDetail.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/Support/LLVM.h"
 
 using namespace mlir;
 using namespace linalg;
 
 //===----------------------------------------------------------------------===//

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

 
 #include "mlir/Dialect/Linalg/Transforms/HoistPadding.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Matchers.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"

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 "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/IR/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 "llvm/ADT/ScopeExit.h"

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

 #include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "llvm/Support/CommandLine.h"
 
 using namespace mlir;

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

 
 #include <type_traits>
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 #include "llvm/ADT/DenseSet.h"
if (m == 0) {
         result = rewriter.create<vector::InsertOp>(loc, bcst, result, d);
       }
     }
     rewriter.replaceOp(op, result);
     return success();
   }
 };
 
-/// Return the number of leftmost dimensions from the first rightmost transposed
-/// dimension found in 'transpose'.
-size_t getNumDimsFromFirstTransposedDim(ArrayRef<int64_t> transpose) {
+/// Given a 'transpose' pattern, prune the rightmost dimensions that are not
+/// transposed.
+void pruneNonTransposedDims(ArrayRef<int64_t> transpose,
+                            SmallVectorImpl<int64_t> &result) {
   size_t numTransposedDims = transpose.size();
   for (size_t transpDim : llvm::reverse(transpose)) {
     if (transpDim != numTransposedDims - 1)
       break;
     numTransposedDims--;
   }
-  return numTransposedDims;
+
+  result.append(transpose.begin(), transpose.begin() + numTransposedDims);
 }
 
 /// Progressive lowering of TransposeOp.
 /// One:
 ///   %x = vector.transpose %y, [1, 0]
 /// is replaced by:
 ///   %z = arith.constant dense<0.000000e+00>
 ///   %0 = vector.extract %y[0, 0]
 ///   %1 = vector.insert %0, %z [0, 0]
 ///   ..
 ///   %x = vector.insert .., .. [.., ..]
 class TransposeOpLowering : public OpRewritePattern<vector::TransposeOp> {
 public:
   using OpRewritePattern<vector::TransposeOp>::OpRewritePattern;
 
   TransposeOpLowering(vector::VectorTransformsOptions vectorTransformOptions,
                       MLIRContext *context)
       : OpRewritePattern<vector::TransposeOp>(context),
         vectorTransformOptions(vectorTransformOptions) {}
 
   LogicalResult matchAndRewrite(vector::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
+    Value input = op.vector();
+    VectorType inputType = op.getVectorType();
     VectorType resType = op.getResultType();
 
     // Set up convenience transposition table.
     SmallVector<int64_t, 4> transp;
     for (auto attr : op.transp())
       transp.push_back(attr.cast<IntegerAttr>().getInt());
 
     if (vectorTransformOptions.vectorTransposeLowering ==
             vector::VectorTransposeLowering::Shuffle &&
         resType.getRank() == 2 && transp[0] == 1 && transp[1] == 0)
       return rewriter.notifyMatchFailure(
           op, "Options specifies lowering to shuffle");
 
     // Handle a true 2-D matrix transpose differently when requested.
     if (vectorTransformOptions.vectorTransposeLowering ==
             vector::VectorTransposeLowering::Flat &&
         resType.getRank() == 2 && transp[0] == 1 && transp[1] == 0) {
       Type flattenedType =
           VectorType::get(resType.getNumElements(), resType.getElementType());
       auto matrix =
-          rewriter.create<vector::ShapeCastOp>(loc, flattenedType, op.vector());
+          rewriter.create<vector::ShapeCastOp>(loc, flattenedType, input);
       auto rows = rewriter.getI32IntegerAttr(resType.getShape()[0]);
       auto columns = rewriter.getI32IntegerAttr(resType.getShape()[1]);
       Value trans = rewriter.create<vector::FlatTransposeOp>(
           loc, flattenedType, matrix, rows, columns);
       rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, resType, trans);
       return success();
     }
 
     // Generate unrolled extract/insert ops. We do not unroll the rightmost
     // (i.e., highest-order) dimensions that are not transposed and leave them
-    // in vector form to improve performance.
-    size_t numLeftmostTransposedDims = getNumDimsFromFirstTransposedDim(transp);
-
-    // The type of the extract operation will be scalar if all the dimensions
-    // are unrolled. Otherwise, it will be a vector with the shape of the
-    // dimensions that are not transposed.
-    Type extractType =
-        numLeftmostTransposedDims == transp.size()
-            ? resType.getElementType()
-            : VectorType::Builder(resType).setShape(
-                  resType.getShape().drop_front(numLeftmostTransposedDims));
-
+    // in vector form to improve performance. Therefore, we prune those
+    // dimensions from the shape/transpose data structures used to generate the
+    // extract/insert ops.
+    SmallVector<int64_t, 4> prunedTransp;
+    pruneNonTransposedDims(transp, prunedTransp);
+    size_t numPrunedDims = transp.size() - prunedTransp.size();
+    auto prunedInShape = inputType.getShape().drop_back(numPrunedDims);
+    SmallVector<int64_t, 4> ones(prunedInShape.size(), 1);
+    auto prunedInStrides = computeStrides(prunedInShape, ones);
+
+    // Generates the extract/insert operations for every scalar/vector element
+    // of the leftmost transposed dimensions. We traverse every transpose
+    // element using a linearized index that we delinearize to generate the
+    // appropriate indices for the extract/insert operations.
     Value result = rewriter.create<arith::ConstantOp>(
         loc, resType, rewriter.getZeroAttr(resType));
-    SmallVector<int64_t, 4> lhs(numLeftmostTransposedDims, 0);
-    SmallVector<int64_t, 4> rhs(numLeftmostTransposedDims, 0);
-    rewriter.replaceOp(op, expandIndices(loc, resType, extractType, 0,
-                                         numLeftmostTransposedDims, transp, lhs,
-                                         rhs, op.vector(), result, rewriter));
-    return success();
-  }
+    int64_t numTransposedElements = ShapedType::getNumElements(prunedInShape);
 
-private:
-  // Builds the indices arrays for the lhs and rhs. Generates the extract/insert
-  // operations when all the ranks go over the last dimension being transposed.
-  Value expandIndices(Location loc, VectorType resType, Type extractType,
-                      int64_t pos, int64_t numLeftmostTransposedDims,
-                      SmallVector<int64_t, 4> &transp,
-                      SmallVector<int64_t, 4> &lhs,
-                      SmallVector<int64_t, 4> &rhs, Value input, Value result,
-                      PatternRewriter &rewriter) const {
-    if (pos >= numLeftmostTransposedDims) {
-      auto ridx = rewriter.getI64ArrayAttr(rhs);
-      auto lidx = rewriter.getI64ArrayAttr(lhs);
-      Value e =
-          rewriter.create<vector::ExtractOp>(loc, extractType, input, ridx);
-      return rewriter.create<vector::InsertOp>(loc, resType, e, result, lidx);
-    }
-    for (int64_t d = 0, e = resType.getDimSize(pos); d < e; ++d) {
-      lhs[pos] = d;
-      rhs[transp[pos]] = d;
-      result = expandIndices(loc, resType, extractType, pos + 1,
-                             numLeftmostTransposedDims, transp, lhs, rhs, input,
-                             result, rewriter);
+    for (int64_t linearIdx = 0; linearIdx < numTransposedElements;
+         ++linearIdx) {
+      auto extractIdxs = delinearize(prunedInStrides, linearIdx);
+      SmallVector<int64_t, 4> insertIdxs(extractIdxs);
+      applyPermutationToVector(insertIdxs, prunedTransp);
+      Value extractOp =
+          rewriter.create<vector::ExtractOp>(loc, input, extractIdxs);
+      result =
+          rewriter.create<vector::InsertOp>(loc, extractOp, result, insertIdxs);
     }
-    return result;
+
+    rewriter.replaceOp(op, result);
+    return success();
   }
 
+private:
   /// Options to control the vector patterns.
   vector::VectorTransformsOptions vectorTransformOptions;
 };
 
 /// Rewrite a 2-D vector.transpose as a sequence of:
 ///   vector.shape_cast 2D -> 1D
 ///   vector.shuffle
 ///   vector.shape_cast 1D -> 2D
 class TransposeOp2DToShuffleLowering
     : public OpRewritePattern<vector::TransposeOp> {
 public:
   using OpRewritePattern<vector::TransposeOp>::OpRewritePattern;
 
   TransposeOp2DToShuffleLowering(
       vector::VectorTransformsOptions vectorTransformOptions,
       MLIRContext *context)
       : OpRewritePattern<vector::TransposeOp>(context),
         vectorTransformOptions(vectorTransformOptions) {}
 
   LogicalResult matchAndRewrite(vector::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 

mlir/test/Dialect/Vector/vector-transpose-lowering.mlir

 // RUN: mlir-opt %s -test-vector-transpose-lowering=eltwise=1 -split-input-file | FileCheck %s --check-prefix=ELTWISE
 // RUN: mlir-opt %s -test-vector-transpose-lowering=shuffle=1 -split-input-file | FileCheck %s --check-prefix=SHUFFLE
 // RUN: mlir-opt %s -test-vector-transpose-lowering=flat=1 -split-input-file | FileCheck %s --check-prefix=FLAT
 // RUN: mlir-opt %s -test-vector-transpose-lowering=avx2=1 -split-input-file | FileCheck %s --check-prefix=AVX2
 
 // ELTWISE-LABEL: func @transpose23
 // ELTWISE-SAME: %[[A:.*]]: vector<2x3xf32>
 // ELTWISE:      %[[Z:.*]] = arith.constant dense<0.000000e+00> : vector<3x2xf32>
 // ELTWISE:      %[[T0:.*]] = vector.extract %[[A]][0, 0] : vector<2x3xf32>
 // ELTWISE:      %[[T1:.*]] = vector.insert %[[T0]], %[[Z]] [0, 0] : f32 into vector<3x2xf32>
-// ELTWISE:      %[[T2:.*]] = vector.extract %[[A]][1, 0] : vector<2x3xf32>
-// ELTWISE:      %[[T3:.*]] = vector.insert %[[T2]], %[[T1]] [0, 1] : f32 into vector<3x2xf32>
-// ELTWISE:      %[[T4:.*]] = vector.extract %[[A]][0, 1] : vector<2x3xf32>
-// ELTWISE:      %[[T5:.*]] = vector.insert %[[T4]], %[[T3]] [1, 0] : f32 into vector<3x2xf32>
-// ELTWISE:      %[[T6:.*]] = vector.extract %[[A]][1, 1] : vector<2x3xf32>
-// ELTWISE:      %[[T7:.*]] = vector.insert %[[T6]], %[[T5]] [1, 1] : f32 into vector<3x2xf32>
-// ELTWISE:      %[[T8:.*]] = vector.extract %[[A]][0, 2] : vector<2x3xf32>
-// ELTWISE:      %[[T9:.*]] = vector.insert %[[T8]], %[[T7]] [2, 0] : f32 into vector<3x2xf32>
+// ELTWISE:      %[[T2:.*]] = vector.extract %[[A]][0, 1] : vector<2x3xf32>
+// ELTWISE:      %[[T3:.*]] = vector.insert %[[T2]], %[[T1]] [1, 0] : f32 into vector<3x2xf32>
+// ELTWISE:      %[[T4:.*]] = vector.extract %[[A]][0, 2] : vector<2x3xf32>
+// ELTWISE:      %[[T5:.*]] = vector.insert %[[T4]], %[[T3]] [2, 0] : f32 into vector<3x2xf32>
+// ELTWISE:      %[[T6:.*]] = vector.extract %[[A]][1, 0] : vector<2x3xf32>
+// ELTWISE:      %[[T7:.*]] = vector.insert %[[T6]], %[[T5]] [0, 1] : f32 into vector<3x2xf32>
+// ELTWISE:      %[[T8:.*]] = vector.extract %[[A]][1, 1] : vector<2x3xf32>
+// ELTWISE:      %[[T9:.*]] = vector.insert %[[T8]], %[[T7]] [1, 1] : f32 into vector<3x2xf32>
 // ELTWISE:      %[[T10:.*]] = vector.extract %[[A]][1, 2] : vector<2x3xf32>
 // ELTWISE:      %[[T11:.*]] = vector.insert %[[T10]], %[[T9]] [2, 1] : f32 into vector<3x2xf32>
 // ELTWISE:      return %[[T11]] : vector<3x2xf32>
 func @transpose23(%arg0: vector<2x3xf32>) -> vector<3x2xf32> {
   %0 = vector.transpose %arg0, [1, 0] : vector<2x3xf32> to vector<3x2xf32>
   return %0 : vector<3x2xf32>
 }