[mlir][vector] Merge accumulator/result transpose into contract

[antiagainst]

Summary

This commit adds a pattern to merge accumulator and result
vector.transpose ops into vector.contract. This kind of
pattern can be generated for NCHW convolution vectorization,
where we use transposes to convert the 1-D NCW convolution
into NWC during vectorization. Merging the transpose would
mean we can avoid materialize vector extract/insert for
transposes and it makes further vector level transformations
easier.

[ThomasRaoux]

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

1107

I wonder if we want to fail in case the contractOp has multiple users as it would duplicate the contraction op?

[ThomasRaoux]

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

1133

nit: the code could be simplified a bit by checking inversePermutation(accTMap) != resTMap

[antiagainst]

Address comments

[antiagainst]

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

1107

Good point. Done.

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

rewriter.replaceOpWithNewOp<mlir::vector::ContractionOp>(
         rewriter.getArrayAttr(llvm::to_vector(llvm::map_range(
             iteratorTypes, [&](IteratorType t) -> mlir::Attribute {
               return IteratorTypeAttr::get(rewriter.getContext(), t);
             }))));
     return success();
   }
 };
 
-/// Merge TransposeOp into ContractionOp user.
+/// Merge LHS/RHS (A/B) TransposeOp into ContractionOp user.
 /// Ex:
 /// ```
 ///   %0 = vector.transpose %arg0, [2, 0, 1]
 ///     : vector<32x16x8xf32> to vector<8x32x16xf32>
 ///   %1 = vector.contract {indexing_maps = [
 ///         affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
 ///         affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
 ///         affine_map<(d0, d1, d2) -> (d0, d1)>],
 ///    iterator_types = ["parallel", "parallel", "reduction"],
 ///    kind = add} %0, %arg1, %cst_f0
 ///    : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
 /// ```
 /// Gets converted to:
 /// ```
 ///   %1 = vector.contract {indexing_maps = [
 ///         affine_map<(d0, d1, d2) -> (d1, d2, d0)>,
 ///         affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
 ///         affine_map<(d0, d1, d2) -> (d0, d1)>],
 ///    iterator_types = ["parallel", "parallel", "reduction"],
 ///    kind = add} %arg0, %arg1, %cst_f0
 ///    : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
 ///  ```
-struct CombineContractTranspose
+struct CombineContractABTranspose final
     : public OpRewritePattern<vector::ContractionOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                 PatternRewriter &rewriter) const override {
     SmallVector<AffineMap, 4> maps =
         llvm::to_vector<4>(contractOp.getIndexingMapsArray());
     Value lhs = contractOp.getLhs();
     Value rhs = contractOp.getRhs();
     size_t index = 0;
     bool changed = false;
     for (Value *operand : {&lhs, &rhs}) {
       AffineMap &map = maps[index++];
       auto transposeOp = operand->getDefiningOp<vector::TransposeOp>();
       if (!transposeOp)
         continue;
-      SmallVector<int64_t> perm;
-      transposeOp.getTransp(perm);
       AffineMap permutationMap = AffineMap::getPermutationMap(
           extractVector<unsigned>(transposeOp.getTransp()),
           contractOp.getContext());
       map = inversePermutation(permutationMap).compose(map);
       *operand = transposeOp.getVector();
       changed = true;
     }
     if (!changed)
       return failure();
     rewriter.replaceOpWithNewOp<vector::ContractionOp>(
         contractOp, lhs, rhs, contractOp.getAcc(),
         rewriter.getAffineMapArrayAttr(maps), contractOp.getIteratorTypes());
     return success();
   }
 };
 
+/// Merges accumulator and result transposes into contract.
+///
+/// For example:
+/// ```mlir
+/// %accT = vector.transpose %acc, [0, 2, 1]
+///   : vector<2x8x4xf32> to vector<2x4x8xf32>
+/// %contract = vector.contract {
+///   indexing_maps = [
+///     affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
+///     affine_map<(d0, d1, d2, d3) -> (d3, d2)>,
+///     affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+///   ],
+///   iterator_types = ["parallel", "parallel", "parallel", "reduction"],
+///   kind = #vector.kind<add>
+/// } %lhs, %rhs, %accT
+///   : vector<2x4x4xf32>, vector<4x8xf32> into vector<2x4x8xf32>
+/// %0 = vector.transpose %contract, [0, 2, 1]
+///   : vector<2x4x8xf32> to vector<2x8x4>
+/// ```
+/// Becomes:
+/// ```mlir
+/// %0 = vector.contract {
+///   indexing_maps = [
+///     affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
+///     affine_map<(d0, d1, d2, d3) -> (d3, d2)>,
+///     affine_map<(d0, d1, d2, d3) -> (d0, d2, d1)>
+///   ],
+///   iterator_types = ["parallel", "parallel", "parallel", "reduction"],
+///   kind = #vector.kind<add>
+/// } %lhs, %rhs, %acc
+///   : vector<2x4x4xf32>, vector<4x8xf32> into vector<2x8x4xf32>
+/// ```
+struct CombineContractResultTranspose final
+    : public OpRewritePattern<vector::TransposeOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransposeOp resTOp,
+                                PatternRewriter &rewriter) const override {
+    auto contractOp = resTOp.getVector().getDefiningOp<vector::ContractionOp>();
+    if (!contractOp || !contractOp->hasOneUse())
+      return failure();
+
+    auto accTOp = contractOp.getAcc().getDefiningOp<vector::TransposeOp>();
+    if (!accTOp)
+      return failure();
+
+    MLIRContext *context = contractOp.getContext();
+    auto maps = llvm::to_vector<3>(contractOp.getIndexingMapsArray());
+    AffineMap contractMap = maps.back();
+
+    // Accumulator transpose performs f(A) -> B. Contract performs g(C) -> B.
+    // To index into A in contract, we need revert(f)(g(C)) -> A.
+    auto accTMap = AffineMap::getPermutationMap(
+        extractVector<unsigned>(accTOp.getTransp()), context);
+
+    // Contract performs g(C) -> D. Result transpose performs h(D) -> E.
+    // To index into E in contract, we need h(g(C)) -> E.
+    auto resTMap = AffineMap::getPermutationMap(
+        extractVector<unsigned>(resTOp.getTransp()), context);
+    auto combinedResMap = resTMap.compose(contractMap);
+
+    // The accumulator and result share the same indexing map. So they should be
+    // the same to be able to merge. This means combinedResMap is the same as
+    // inversePermutation(accTMap).compose(contractMap), which means
+    if (inversePermutation(accTMap) != resTMap)
+      return failure();
+    maps.back() = combinedResMap;
+
+    rewriter.replaceOpWithNewOp<vector::ContractionOp>(
+        resTOp, contractOp.getLhs(), contractOp.getRhs(), accTOp.getVector(),
+        rewriter.getAffineMapArrayAttr(maps), contractOp.getIteratorTypes());
+    return success();
+  }
+};
+
 /// Merge BroadcastOp into ContractionOp user.
 /// Ex:
 /// ```
 ///   %0 = vector.broadcast %arg0 : vector<32x16xf32> to vector<8x32x16xf32>
 ///   %1 = vector.contract {indexing_maps = [
 ///         affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
 ///         affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
 ///         affine_map<(d0, d1, d2) -> (d0, d1)>],
LogicalResult matchAndRewrite(CastOpInterface op,
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
         op, op->getResult(0).getType(), castOp->getResult(0));
     return success();
   }
 };
 
 /// Reorders elementwise(transpose) to transpose(elementwise). This makes
 /// transpose ops and contraction ops closer, which kicks in
-/// CombineContractTranspose pattern when elementwise ops are between these
+/// CombineContractABTranspose pattern when elementwise ops are between these
 /// operations. Ex:
 /// ```
 /// %at = vector.transpose %a, [1, 0]: vector<4x2xf32> to vector<2x4xf32>
 /// %bt = vector.transpose %b, [1, 0]: vector<4x2xf32> to vector<2x4xf32>
 /// %r = arith.addf %at, %bt : vector<2x4xf32>
 /// ```
 /// Gets converted to:
 /// ```
void mlir::vector::populateVectorTransposeLoweringPatterns(
     PatternBenefit benefit) {
   patterns.add<TransposeOpLowering, TransposeOp2DToShuffleLowering>(
       options, patterns.getContext(), benefit);
 }
 
 void mlir::vector::populateVectorReductionToContractPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {
   patterns.add<MultiReduceToContract, CombineContractBroadcast,
-               CombineContractTranspose, ReorderCastOpsOnBroadcast,
-               ReorderElementwiseOpsOnTranspose>(patterns.getContext(),
-                                                 benefit);
+               CombineContractABTranspose, CombineContractResultTranspose,
+               ReorderCastOpsOnBroadcast, ReorderElementwiseOpsOnTranspose>(
+      patterns.getContext(), benefit);
 }
 
 void mlir::vector::
     populateVectorTransferCollapseInnerMostContiguousDimsPatterns(
         RewritePatternSet &patterns, PatternBenefit benefit) {
   patterns.add<DropInnerMostUnitDims>(patterns.getContext(), benefit);
 }
 

mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir

 //       CHECK:   vector.transpose
 //       CHECK:   arith.addf
 func.func @transpose_elementwise_diff_map(%a : vector<4x6x3x2xf32>, %b: vector<6x2x4x3xf32>) -> vector<6x4x2x3xf32> {
   %at = vector.transpose %a, [1, 0, 3, 2]: vector<4x6x3x2xf32> to vector<6x4x2x3xf32>
   %bt = vector.transpose %b, [0, 2, 1, 3]: vector<6x2x4x3xf32> to vector<6x4x2x3xf32>
   %r = arith.addf %at, %bt : vector<6x4x2x3xf32>
   return %r : vector<6x4x2x3xf32>
 }
+
+// -----
+
+// CHECK-DAG: #[[$LHS_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
+// CHECK-DAG: #[[$RHS_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>
+// CHECK-DAG: #[[$ACC_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1)>
+
+// CHECK-LABEL: func.func @contract_result_transpose
+//  CHECK-SAME: (%[[LHS:.+]]: vector<2x4x4xf32>, %[[RHS:.+]]: vector<4x8xf32>, %[[ACC:.+]]: vector<2x8x4xf32>)
+//       CHECK:   %[[CONTRACT:.+]] = vector.contract
+//  CHECK-SAME:     indexing_maps = [#[[$LHS_MAP]], #[[$RHS_MAP]], #[[$ACC_MAP]]]
+//  CHECK-SAME:     iterator_types = ["parallel", "parallel", "parallel", "reduction"]
+//  CHECK-SAME:     kind = #vector.kind<add>
+//  CHECK-SAME:     %[[LHS]], %[[RHS]], %[[ACC]]
+//       CHECK:   return %[[CONTRACT]]
+func.func @contract_result_transpose(%lhs : vector<2x4x4xf32>, %rhs: vector<4x8xf32>, %acc: vector<2x8x4xf32>) -> vector<2x8x4xf32> {
+  %accT = vector.transpose %acc, [0, 2, 1] : vector<2x8x4xf32> to vector<2x4x8xf32>
+  %contract = vector.contract {
+    indexing_maps = [
+      affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
+      affine_map<(d0, d1, d2, d3) -> (d3, d2)>,
+      affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+    ],
+    iterator_types = ["parallel", "parallel", "parallel", "reduction"],
+    kind = #vector.kind<add>
+  } %lhs, %rhs, %accT : vector<2x4x4xf32>, vector<4x8xf32> into vector<2x4x8xf32>
+  %resT = vector.transpose %contract, [0, 2, 1] : vector<2x4x8xf32> to vector<2x8x4xf32>
+  return %resT : vector<2x8x4xf32>
+}