[mlir][Vector] NFC - Use matchAndRewrite in ContractionOp lowering patterns

[nicolasvasilache]

Summary

Replace the use of separate match and rewrite which unnecessarily duplicates logic.

[bkramer]

lgtm

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

public:
   }
 
   ContractionOpToMatmulOpLowering(
       vector::VectorTransformsOptions vectorTransformsOptions,
       MLIRContext *context, FilterConstraintType constraint = defaultFilter)
       : OpRewritePattern<vector::ContractionOp>(context),
         vectorTransformsOptions(vectorTransformsOptions), filter(constraint) {}
 
-  LogicalResult match(vector::ContractionOp op) const override;
-  void rewrite(vector::ContractionOp op,
+  LogicalResult matchAndRewrite(vector::ContractionOp op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   /// Options to control the vector patterns.
   vector::VectorTransformsOptions vectorTransformsOptions;
   FilterConstraintType filter;
 };
 
 /// Progressive lowering of a `vector.contract %a, %b, %c` with row-major matmul
public:
   }
 
   ContractionOpToOuterProductOpLowering(
       vector::VectorTransformsOptions vectorTransformsOptions,
       MLIRContext *context, FilterConstraintType constraint = defaultFilter)
       : OpRewritePattern<vector::ContractionOp>(context),
         vectorTransformsOptions(vectorTransformsOptions), filter(constraint) {}
 
-  LogicalResult match(vector::ContractionOp op) const override;
-  void rewrite(vector::ContractionOp op,
+  LogicalResult matchAndRewrite(vector::ContractionOp op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   /// Options to control the vector patterns.
   vector::VectorTransformsOptions vectorTransformsOptions;
   FilterConstraintType filter;
 };
 
 /// Progressive lowering of ContractionOp.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 ///    %flattened_d = vector.matmul %flattened_a, %flattened_b
 ///    %d = vector.shape_cast %%flattened_d
 ///    %e = add %c, %d
 /// ```
 /// `vector.matmul` later lowers to `llvm.matrix.multiply`.
 //
 /// This only kicks in when VectorTransformsOptions is set to OuterProduct and
 /// the vector.contract op is a row-major matrix multiply.
-LogicalResult
-ContractionOpToMatmulOpLowering::match(vector::ContractionOp op) const {
+LogicalResult ContractionOpToMatmulOpLowering::matchAndRewrite(
+    vector::ContractionOp op, PatternRewriter &rewriter) const {
   // TODO: implement masks
   if (llvm::size(op.masks()) != 0)
     return failure();
-
   if (vectorTransformsOptions.vectorContractLowering !=
       vector::VectorContractLowering::Matmul)
     return failure();
-
   if (failed(filter(op)))
     return failure();
 
   auto iteratorTypes = op.iterator_types().getValue();
   if (!isParallelIterator(iteratorTypes[0]) ||
       !isParallelIterator(iteratorTypes[1]) ||
       !isReductionIterator(iteratorTypes[2]))
     return failure();
 
   if (!isRowMajorMatmul(op.indexing_maps()))
     return failure();
 
-  return success();
-}
+  Type elementType = op.getLhsType().getElementType();
+  if (!elementType.isIntOrFloat())
+    return failure();
 
-void ContractionOpToMatmulOpLowering::rewrite(vector::ContractionOp op,
-                                              PatternRewriter &rewriter) const {
   VectorType lhsType = op.getLhsType();
   VectorType rhsType = op.getRhsType();
   int64_t lhsRows = lhsType.getDimSize(0);
   int64_t lhsColumns = lhsType.getDimSize(1);
   int64_t rhsColumns = rhsType.getDimSize(1);
 
   Type flattenedLHSType =
       VectorType::get(lhsType.getNumElements(), lhsType.getElementType());
   Type flattenedRHSType =
       VectorType::get(rhsType.getNumElements(), rhsType.getElementType());
   auto lhs = rewriter.create<vector::ShapeCastOp>(op.getLoc(), flattenedLHSType,
                                                   op.lhs());
   auto rhs = rewriter.create<vector::ShapeCastOp>(op.getLoc(), flattenedRHSType,
                                                   op.rhs());
 
   Value mul = rewriter.create<vector::MatmulOp>(op.getLoc(), lhs, rhs, lhsRows,
                                                 lhsColumns, rhsColumns);
   mul = rewriter.create<vector::ShapeCastOp>(op.getLoc(), op.acc().getType(),
                                              mul);
-  Type elementType = op.getLhsType().getElementType();
-  assert(elementType.isIntOrFloat());
   if (elementType.isa<IntegerType>())
     rewriter.replaceOpWithNewOp<AddIOp>(op, op.acc(), mul);
   else
     rewriter.replaceOpWithNewOp<AddFOp>(op, op.acc(), mul);
+
+  return success();
 }
 
 /// Progressively lower a `vector.contract %a, %b, %c` with row-major matmul
 /// semantics to a reduction_size-unrolled sequence:
 /// ```
 ///    %at = vector.transpose %a, [1, 0]
 ///    %bRow0 = vector.extract %b[0]
 ///    %atRow0 = vector.extract %at[0]
 ///    %c0 = vector.outerproduct %atRow0, %bRow0, %c
 ///    ...
 ///    %bRowK = vector.extract %b[K]
 ///    %atRowK = vector.extract %at[K]
 ///    %cK = vector.outerproduct %atRowK, %bRowK, %cK-1
 /// ```
 ///
 /// This only kicks in when VectorTransformsOptions is set to OuterProduct but
 /// otherwise supports any layout permutation of the matrix-multiply.
-LogicalResult
-ContractionOpToOuterProductOpLowering::match(vector::ContractionOp op) const {
+LogicalResult ContractionOpToOuterProductOpLowering::matchAndRewrite(
+    vector::ContractionOp op, PatternRewriter &rewriter) const {
   // TODO: implement masks
   if (llvm::size(op.masks()) != 0)
     return failure();
 
   if (vectorTransformsOptions.vectorContractLowering !=
       vector::VectorContractLowering::OuterProduct)
     return failure();
 
   if (failed(filter(op)))
     return failure();
 
-  // Determine if the parallel/reduction structure matches something
-  // that can be expressed a reduction_size unrolled sequence.
-  using MapList = ArrayRef<ArrayRef<AffineExpr>>;
-  auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
-  AffineExpr m, n, k;
-  bindDims(op.getContext(), m, n, k);
-  auto iteratorTypes = op.iterator_types().getValue();
-  SmallVector<AffineMap, 4> maps = op.getIndexingMaps();
-  if (isParallelIterator(iteratorTypes[0]) &&
-      isParallelIterator(iteratorTypes[1]) &&
-      isReductionIterator(iteratorTypes[2])) {
-    //
-    // Two outer parallel, one inner reduction (matmat flavor).
-    // When lowering to outerproduct we can support all permutations.
-    //
-    if (maps != infer({{m, k}, {k, n}, {m, n}}) &&
-        maps != infer({{m, k}, {n, k}, {m, n}}) &&
-        maps != infer({{k, m}, {k, n}, {m, n}}) &&
-        maps != infer({{k, m}, {n, k}, {m, n}}) &&
-        maps != infer({{m, k}, {k, n}, {n, m}}) &&
-        maps != infer({{m, k}, {n, k}, {n, m}}) &&
-        maps != infer({{k, m}, {k, n}, {n, m}}) &&
-        maps != infer({{k, m}, {n, k}, {n, m}}))
-      return failure();
-    return success();
-  } else if (isParallelIterator(iteratorTypes[0]) &&
-             isReductionIterator(iteratorTypes[1])) {
-    //
-    // One outer parallel, one inner reduction (matvec flavor)
-    // See if a series of AXPY operations chained through FMA operations
-    // could replace the default DOT implementation.
-    //
-    if (maps != infer({{m, n}, {n}, {m}}) && // mat-vec
-        maps != infer({{n, m}, {n}, {m}}) && // mat-trans-vec
-        maps != infer({{n}, {m, n}, {m}}) && // vec-mat
-        maps != infer({{n}, {n, m}, {m}}))   // vec-mat-trans
-      return failure();
-    return success();
-  }
-  return failure();
-}
-
-void ContractionOpToOuterProductOpLowering::rewrite(
-    vector::ContractionOp op, PatternRewriter &rewriter) const {
   Location loc = op.getLoc();
   int64_t reductionSize = 0;
   VectorType lhsType = op.getLhsType();
   Value lhs = op.lhs(), rhs = op.rhs(), res = op.acc();
 
   // Set up the parallel/reduction structure in right form.
   using MapList = ArrayRef<ArrayRef<AffineExpr>>;
   auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
if (maps == infer({{m, k}, {k, n}, {m, n}})) {
       reductionSize = lhsType.getDimSize(0);
       std::swap(lhs, rhs);
     } else if (maps == infer({{k, m}, {n, k}, {n, m}})) {
       // Just permute the rhs.
       reductionSize = lhsType.getDimSize(0);
       Value tmp = lhs;
       lhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
       rhs = tmp;
-    }
     } else {
+      return failure();
+    }
+  } else if (isParallelIterator(iteratorTypes[0]) &&
+             isReductionIterator(iteratorTypes[1])) {
     //
     // One outer parallel, one inner reduction (matvec flavor)
     //
-    assert(isParallelIterator(iteratorTypes[0]) &&
-           isReductionIterator(iteratorTypes[1]));
     if (maps == infer({{m, n}, {n}, {m}})) {
       // Case mat-vec: transpose.
       reductionSize = lhsType.getDimSize(1);
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
     } else if (maps == infer({{n, m}, {n}, {m}})) {
       // Case mat-trans-vec: ready to go.
       reductionSize = lhsType.getDimSize(0);
     } else if (maps == infer({{n}, {m, n}, {m}})) {
       // Case vec-mat: swap and transpose.
       reductionSize = lhsType.getDimSize(0);
       std::swap(lhs, rhs);
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
     } else if (maps == infer({{n}, {n, m}, {m}})) {
       // Case vec-mat-trans: swap and ready to go.
       reductionSize = lhsType.getDimSize(0);
       std::swap(lhs, rhs);
+    } else {
+      return failure();
     }
+  } else {
+    return failure();
   }
   assert(reductionSize > 0);
 
   // Unroll outer-products along reduction.
   for (int64_t k = 0; k < reductionSize; ++k) {
     Value a = rewriter.create<vector::ExtractOp>(op.getLoc(), lhs, k);
     Value b = rewriter.create<vector::ExtractOp>(op.getLoc(), rhs, k);
     res = rewriter.create<vector::OuterProductOp>(op.getLoc(), a, b, res);
   }
   rewriter.replaceOp(op, res);
+  return success();
 }
 
 /// Progressive lowering of ContractionOp.
 /// One:
 ///   %x = vector.contract with at least one free/batch dimension
 /// is replaced by:
 ///   %a = vector.contract with one less free/batch dimension
 ///   %b = vector.contract with one less free/batch dimension
 ///   ..
 ///   %x = combine %a %b ..
 /// until a pure contraction is reached (no free/batch dimensions),
 /// which is replaced by a dot-product.
 ///
 /// This only kicks in when either VectorTransformsOptions is set
 /// to DOT or when other contraction patterns fail.
 //
 // TODO: break down into transpose/reshape/cast ops
 //               when they become available to avoid code dup
 // TODO: investigate lowering order impact on performance
 LogicalResult
 ContractionOpLowering::matchAndRewrite(vector::ContractionOp op,
                                        PatternRewriter &rewriter) const {
-
   // TODO: implement masks.
   if (llvm::size(op.masks()) != 0)
     return failure();
 
   if (failed(filter(op)))
     return failure();
 
   // TODO: support mixed mode contract lowering.
   if (op.getLhsType().getElementType() !=
           getElementTypeOrSelf(op.getAccType()) ||
       op.getRhsType().getElementType() != getElementTypeOrSelf(op.getAccType()))
     return failure();
 
   // TODO: implement benefits, cost models.
   MLIRContext *ctx = op.getContext();
   ContractionOpToMatmulOpLowering pat1(vectorTransformsOptions, ctx);
-  if (succeeded(pat1.match(op)))
-    return failure();
+  if (succeeded(pat1.matchAndRewrite(op, rewriter)))
+    return success();
   ContractionOpToOuterProductOpLowering pat2(vectorTransformsOptions, ctx);
-  if (succeeded(pat2.match(op)))
-    return failure();
+  if (succeeded(pat2.matchAndRewrite(op, rewriter)))
+    return success();
 
   // Find first batch dimension in LHS/RHS, and lower when found.
   std::vector<std::pair<int64_t, int64_t>> batchDimMap = op.getBatchDimMap();
   if (!batchDimMap.empty()) {
     int64_t lhsIndex = batchDimMap[0].first;
     int64_t rhsIndex = batchDimMap[0].second;
     rewriter.replaceOp(op, lowerParallel(op, lhsIndex, rhsIndex, rewriter));
     return success();