[mlir][vector] Fix CombineContractBroadcast for invalid cases

[antiagainst]

Summary

There are cases that after combining vector.contract and
vector.broadcast, the generated vector.contract's operands
do not have parallel or reduction pair in LHS and RHS at all.
Such cases may fail vector.contract verification. Explicitly
guard against such cases.

[antiagainst]

Update

[antiagainst]

Update

[ThomasRaoux]

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

1058

this is only correct for float and add kind of contract?

[antiagainst]

Address comments

[antiagainst]

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

1058

Ah, great catch! Fixed.

[antiagainst]

Update to fix another invalid case

[ThomasRaoux]

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

1066–1076

I don't think we can really remove the reduction dimension unless they are unit dimension?
For instance:

func.func @contract_broadcast_fma(%a: vector<4xf32>, %b: vector<4xf32>, %c: vector<4xf32>) -> vector<4xf32> {
  %bcast_a = vector.broadcast %a : vector<4xf32> to vector<1x2x4xf32>
  %bcast_b = vector.broadcast %b : vector<4xf32> to vector<1x2x4xf32>
  %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>
  } %bcast_a, %bcast_b, %c : vector<1x2x4xf32>, vector<1x2x4xf32> into vector<4xf32>
  return %contract: vector<4xf32>
}

shouldn't become:

func.func @contract_broadcast_fma(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: vector<4xf32>) -> vector<4xf32> {
  %0 = vector.fma %arg0, %arg1, %arg2 : vector<4xf32>
  return %0 : vector<4xf32>
}

1069–1074

I find it a bit odd that we have lowering to fma in the pattern that tries to combine contract and broadcast, why can't this be a separate pattern that would be called by user if it wants this kind of lowering.

[antiagainst]

Remove vector.fma generation

[antiagainst]

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

1066–1076

You are right.. This is problematic. Removed now.

1069–1074

Makes sense. Removed this part now.

[antiagainst]

Fix func.func

[ThomasRaoux]

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

1051–1060

nit: can we reverse the condition to do early exit like the rest of the function?

1053

I don't understand why we need the second condition? As long as we have a reduction being used it should be okay to generate the new contraction op?

[antiagainst]

Address comments

[antiagainst]

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

1051–1060

Done.

1053

It's required by contraction op verification. For full reduction, valid contraction op expects scalar accumulator/result https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Vector/IR/VectorOps.cpp#L674-L676. So we need to make sure we are checking that here. The second test case checks this.

[nicolasvasilache]

would be nice to refactor to directly call into a new verifyContractionOpImpl helper in the same way that we have tensor::verifyInsertSliceOp

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

if (!changed)
       return failure();
     rewriter.replaceOpWithNewOp<vector::ContractionOp>(
         contractOp, lhs, rhs, contractOp.getAcc(),
         rewriter.getAffineMapArrayAttr(maps), contractOp.getIteratorTypes());
     return success();
   }
 };
 
+/// Returns true if the given `lhsMap` and `rhsMap` from a vector.contract op
+/// has a reduction dimension access pair.
+static bool hasReductionPair(AffineMap lhsMap, AffineMap rhsMap,
+                             ArrayAttr iteratorTypes) {
+  for (const auto &it : llvm::enumerate(iteratorTypes)) {
+    if (!isReductionIterator(it.value()))
+      continue;
+    auto lhsDim = getResultIndex(lhsMap, it.index());
+    auto rhsDim = getResultIndex(rhsMap, it.index());
+    if (lhsDim && rhsDim)
+      return true;
+  }
+  return false;
+}
+
 /// 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)>],
for (Value *operand : {&lhs, &rhs}) {
           AffineMap::get(broadcast.getVectorType().getRank(), 0, originalDims,
                          contractOp.getContext());
       map = broadcastMap.compose(map);
       *operand = broadcast.getSource();
       changed = true;
     }
     if (!changed)
       return failure();
+
+    ArrayAttr iteratorTypes = contractOp.getIteratorTypes();
+    // We need to make sure at least one reduction dimension is actually used to
+    // generate valid vector.contract ops.
+    if (hasReductionPair(maps[0], maps[1], iteratorTypes)) {
       rewriter.replaceOpWithNewOp<vector::ContractionOp>(
           contractOp, lhs, rhs, contractOp.getAcc(),
           rewriter.getAffineMapArrayAttr(maps), contractOp.getIteratorTypes());
+    } else {
+      // After combining, all access are through parallel dimensions. It can be
+      // simplified into a vector.fma op if all maps are the same.
+      if (!llvm::is_splat(maps))
+        return failure();
+      rewriter.replaceOpWithNewOp<vector::FMAOp>(contractOp, lhs, rhs,
+                                                 contractOp.getAcc());
+    }
     return success();
   }
 };
 
 /// Reorders cast(broadcast) to broadcast(cast). This makes broadcast ops and
 /// contraction ops closer, which kicks in CombineContractBroadcast pattern when
 /// casting ops are around these operations.
 /// Ex:
 /// ```
 ///   %0 = vector.broadcast %arg0 : vector<32x16xi8> to vector<8x32x16xi8>
 ///   %1 = arith.extsi %0 : vector<8x32x16xi8> to vector<8x32x16xi32>
 /// ```
 /// Gets converted to:
 /// ```
 ///   %0 = arith.extsi %0 : vector<32x16xi8> to vector<32x16xi32>
 ///   %1 = vector.broadcast %arg0 : vector<32x16xi32> to vector<8x32x16xi32>
 /// ```
 struct ReorderCastOpsOnBroadcast
     : public OpInterfaceRewritePattern<CastOpInterface> {
   using OpInterfaceRewritePattern<CastOpInterface>::OpInterfaceRewritePattern;
 
   LogicalResult matchAndRewrite(CastOpInterface op,
                                 PatternRewriter &rewriter) const override {
     if (op->getNumOperands() != 1)
       return failure();
     auto bcastOp = op->getOperand(0).getDefiningOp<vector::BroadcastOp>();
     if (!bcastOp)
       return failure();
 
     Type castResTy = getElementTypeOrSelf(op->getResult(0));
     if (auto vecTy = bcastOp.getSourceType().dyn_cast<VectorType>())
       castResTy = VectorType::get(vecTy.getShape(), castResTy);
     auto castOp =
         rewriter.create(op->getLoc(), op->getName().getIdentifier(),
                         bcastOp.getSource(), castResTy, op->getAttrs());
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
         op, op->getResult(0).getType(), castOp->getResult(0));
     return success();
   }
 };
 
 /// Reorders cast(transpose) to transpose(cast). This makes broadcast ops and
 /// contraction ops closer, which kicks in CombineContractTranspose pattern when
 /// casting ops are around these operations.
 /// Ex:
 /// ```
 ///   %0 = vector.transpose %arg0, [2, 0, 1]

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

func @contract_broadcast(
   %cst = arith.constant dense<0.000000e+00> : vector<8x32xf32>
   %0 = vector.broadcast %arg0 : vector<32x16xf32> to vector<8x32x16xf32>
   %1 = vector.contract {indexing_maps = [#map0, #map0, #map1],
     iterator_types = ["parallel", "parallel", "reduction"],
     kind = #vector.kind<add>} %0, %arg1, %cst : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
   return %1 : vector<8x32xf32>
 }
 
+// -----
+
+// CHECK-LABEL: contract_broadcast_fma
+//  CHECK-SAME: (%[[A:.+]]: vector<4xf32>, %[[B:.+]]: vector<4xf32>, %[[C:.+]]: vector<4xf32>)
+//       CHECK:   %[[FMA:.+]] = vector.fma %[[A]], %[[B]], %[[C]] : vector<4xf32>
+//       CHECK:   return %[[FMA]] : vector<4xf32
+func @contract_broadcast_fma(%a: vector<4xf32>, %b: vector<4xf32>, %c: vector<4xf32>) -> vector<4xf32> {
+  %bcast_a = vector.broadcast %a : vector<4xf32> to vector<1x1x4xf32>
+  %bcast_b = vector.broadcast %b : vector<4xf32> to vector<1x1x4xf32>
+  %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>
+  } %bcast_a, %bcast_b, %c : vector<1x1x4xf32>, vector<1x1x4xf32> into vector<4xf32>
+  return %contract: vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @contract_broadcast_no_reduction_pair
+//       CHECK: vector.broadcast
+//       CHECK: vector.contract
+func @contract_broadcast_no_reduction_pair(%a: vector<1xf32>, %b: vector<4xf32>, %c: vector<4xf32>) -> vector<4xf32> {
+  %bcast = vector.broadcast %b : vector<4xf32> to vector<1x4xf32>
+  %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, %bcast, %c : vector<1xf32>, vector<1x4xf32> into vector<4xf32>
+  return %contract: vector<4xf32>
+}
+
 //===----------------------------------------------------------------------===//
 // Reorder casting ops and vector ops. The casting ops have almost identical
 // pattern, so only arith.extsi op is tested.
 //===----------------------------------------------------------------------===//
 
 // -----
 
 func @broadcast_vector_extsi(%a : vector<4xi8>) -> vector<2x4xi32> {