[mlir][vector] Canonicalize contract to fma when possible

[antiagainst]

Summary

If all reduction dimensions are of size one for a vector.contract
op, we might be able to canonicalize it into a plain vector.fma
op. This allows us to clean up the IR earlier and generate large
vector.fma ops before lowering, which may unroll the parell
dimensions. By then it might be too late.

[antiagainst]

Fix return value

[antiagainst]

Not needed anymore.

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/bit.h"
 #include <numeric>
 
 #include "mlir/Dialect/Vector/IR/VectorOpsDialect.cpp.inc"
 // Pull in all enum type and utility function definitions.
 #include "mlir/Dialect/Vector/IR/VectorOpsEnums.cpp.inc"
 
LogicalResult matchAndRewrite(AddOpType addOp,
 
     Value a = addOp->getOperand(0), b = addOp->getOperand(1);
     vector::ContractionOp contract = canonicalize(a, b);
     contract = contract ? contract : canonicalize(b, a);
     return contract ? success() : failure();
   }
 };
 
+/// Canonicalizes vector.contract with all size one reduction dimensions to
+/// vector.fma ops when possible.
+struct ConvertContractToFMA : public OpRewritePattern<vector::ContractionOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
+                                PatternRewriter &rewriter) const override {
+    // The reduction part should be addition on float element types.
+    if (contractOp.getKind() != vector::CombiningKind::ADD ||
+        !getElementTypeOrSelf(contractOp.getType()).isa<FloatType>())
+      return failure();
+
+    ArrayRef<int64_t> lhsShape = contractOp.getLhsType().getShape();
+    ArrayRef<int64_t> rhsShape = contractOp.getRhsType().getShape();
+
+    // Collect reduction dimensions for both LHS and RHS.
+    auto reductionDimPairs = contractOp.getContractingDimMap();
+    SmallVector<unsigned> lhsReductionDims, rhsReductionDims;
+    lhsReductionDims.reserve(reductionDimPairs.size());
+    rhsReductionDims.reserve(reductionDimPairs.size());
+    for (const auto &lhsRhs : reductionDimPairs) {
+      // Require all reduction dimensions to be of size one to canonicalize.
+      if (lhsShape[lhsRhs.first] != 1 || rhsShape[lhsRhs.second] != 1)
+        return failure();
+      lhsReductionDims.push_back(lhsRhs.first);
+      rhsReductionDims.push_back(lhsRhs.second);
+    }
+
+    // Make sure all reduction dimensions are leading ones, so that we can
+    // generate vector.extract ops later to perform dimension reduction.
+    llvm::sort(lhsReductionDims);
+    llvm::sort(rhsReductionDims);
+    for (const auto &kv : llvm::enumerate(lhsReductionDims))
+      if (kv.index() != kv.value())
+        return failure();
+    for (const auto &kv : llvm::enumerate(rhsReductionDims))
+      if (kv.index() != kv.value())
+        return failure();
+
+    // Filter out reduction dimensions for LHS and RHS.
+    SmallVector<unsigned> lhsParallelDims, rhsParallelDims;
+    SmallVector<int64_t> newLhsShape, newRhsShape;
+    for (unsigned i = 0, e = lhsShape.size(); i < e; ++i)
+      if (!llvm::is_contained(lhsReductionDims, i)) {
+        lhsParallelDims.push_back(i);
+        newLhsShape.push_back(lhsShape[i]);
+      }
+    for (unsigned i = 0, e = rhsShape.size(); i < e; ++i)
+      if (!llvm::is_contained(rhsReductionDims, i)) {
+        rhsParallelDims.push_back(i);
+        newRhsShape.push_back(rhsShape[i]);
+      }
+
+    auto newLhsMap = contractOp.getIndexingMaps()[0].getSubMap(lhsParallelDims);
+    auto newRhsMap = contractOp.getIndexingMaps()[1].getSubMap(rhsParallelDims);
+    auto accMap = contractOp.getIndexingMaps()[2];
+    // After removing all size one reduction dimensions, all indexing maps
+    // should be the same to canonicalize.
+    if (newLhsMap != accMap || newRhsMap != accMap)
+      return failure();
+
+    Location loc = contractOp.getLoc();
+    auto newLhsType =
+        VectorType::get(newLhsShape, contractOp.getLhsType().getElementType());
+    auto newRhsType =
+        VectorType::get(newRhsShape, contractOp.getRhsType().getElementType());
+    SmallVector<int64_t, 4> lhsOffsets(lhsReductionDims.size(), 0);
+    SmallVector<int64_t, 4> rhsOffsets(rhsReductionDims.size(), 0);
+    Value newLhs = rewriter.create<vector::ExtractOp>(
+        loc, newLhsType, contractOp.getLhs(),
+        rewriter.getI64ArrayAttr(lhsOffsets));
+    Value newRhs = rewriter.create<vector::ExtractOp>(
+        loc, newRhsType, contractOp.getRhs(),
+        rewriter.getI64ArrayAttr(rhsOffsets));
+
+    rewriter.replaceOpWithNewOp<vector::FMAOp>(contractOp, newLhs, newRhs,
+                                               contractOp.getAcc());
+    return success();
+  }
+};
+
 void ContractionOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
   results.add<CanonicalizeContractAdd<arith::AddIOp>,
-              CanonicalizeContractAdd<arith::AddFOp>>(context);
+              CanonicalizeContractAdd<arith::AddFOp>, ConvertContractToFMA>(
+      context);
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
 void vector::ExtractElementOp::build(OpBuilder &builder, OperationState &result,
                                      Value source) {

mlir/test/Dialect/Vector/canonicalize.mlir

 // -----
 
 // CHECK-LABEL: func @dont_reduce_one_element_vector
 //       CHECK: vector.reduction
 func.func @dont_reduce_one_element_vector(%a : vector<4xf32>) -> f32 {
   %s = vector.reduction <add>, %a : vector<4xf32> into f32
   return %s : f32
 }
+
+// -----
+
+// CHECK-LABEL: func @contract_to_fma
+//  CHECK-SAME: (%[[A:.+]]: vector<1x1x4xf32>, %[[B:.+]]: vector<1x1x4xf32>, %[[C:.+]]: vector<4xf32>)
+//       CHECK:   %[[EXTRACTA:.+]] = vector.extract %[[A]][0, 0] : vector<1x1x4xf32>
+//       CHECK:   %[[EXTRACTB:.+]] = vector.extract %[[B]][0, 0] : vector<1x1x4xf32>
+//       CHECK:   %[[FMA:.+]] = vector.fma %[[EXTRACTA]], %[[EXTRACTB]], %[[C]] : vector<4xf32>
+//       CHECK:   return %[[FMA]] : vector<4xf32>
+func.func @contract_to_fma(%a: vector<1x1x4xf32>, %b: vector<1x1x4xf32>, %c: vector<4xf32>) -> vector<4xf32> {
+  %contract = vector.contract {
+    indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d0)>],
+    iterator_types = ["parallel", "reduction", "reduction"], kind = #vector.kind<add>
+  } %a, %b, %c : vector<1x1x4xf32>, vector<1x1x4xf32> into vector<4xf32>
+  return %contract: vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @contract_to_fma_i32
+//       CHECK: vector.contract
+func.func @contract_to_fma_i32(%a: vector<1x1x4xi32>, %b: vector<1x1x4xi32>, %c: vector<4xi32>) -> vector<4xi32> {
+  %contract = vector.contract {
+    indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d0)>],
+    iterator_types = ["parallel", "reduction", "reduction"], kind = #vector.kind<add>
+  } %a, %b, %c : vector<1x1x4xi32>, vector<1x1x4xi32> into vector<4xi32>
+  return %contract: vector<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @contract_to_fma_maxf
+//       CHECK: vector.contract
+func.func @contract_to_fma_maxf(%a: vector<1x1x4xf32>, %b: vector<1x1x4xf32>, %c: vector<4xf32>) -> vector<4xf32> {
+  %contract = vector.contract {
+    indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d0)>],
+    iterator_types = ["parallel", "reduction", "reduction"], kind = #vector.kind<maxf>
+  } %a, %b, %c : vector<1x1x4xf32>, vector<1x1x4xf32> into vector<4xf32>
+  return %contract: vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @contract_to_fma_diff_map
+//       CHECK: vector.contract
+func.func @contract_to_fma_diff_map(%a: vector<1xf32>, %b: vector<1x4xf32>, %c: vector<4xf32>) -> vector<4xf32> {
+  %contract = vector.contract {
+    indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0)>],
+    iterator_types = ["parallel", "reduction"], kind = #vector.kind<add>
+  } %a, %b, %c : vector<1xf32>, vector<1x4xf32> into vector<4xf32>
+  return %contract: vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @contract_to_fma_non1_reduction
+//       CHECK: vector.contract
+func.func @contract_to_fma_non1_reduction(%a: vector<1x2x4xf32>, %b: vector<1x2x4xf32>, %c: vector<4xf32>) -> vector<4xf32> {
+  %contract = vector.contract {
+    indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d0)>],
+    iterator_types = ["parallel", "reduction", "reduction"], kind = #vector.kind<add>
+  } %a, %b, %c : vector<1x2x4xf32>, vector<1x2x4xf32> into vector<4xf32>
+  return %contract: vector<4xf32>
+}