CombineContractBroadcast should not create dims unused in LHS+RHS

[nicolasvasilache]

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

1175

Please add (i.e. the ContractionOp verifier fails).

1176

Can we factor out the common verification bits that are useful here and reuse with the verifier too?
This has a flavor similar to verifyInsertSliceOp in TensorOps.cpp and that has served us well in the past.

[Benoit]

Use also in verifier, and in one place in AffineMap.cpp (that this logic was borrowed from in the first place)

mlir/include/mlir/IR/AffineMap.h

 
 #ifndef MLIR_IR_AFFINEMAP_H
 #define MLIR_IR_AFFINEMAP_H
 
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/SmallBitVector.h"
 
 namespace llvm {
 class SmallBitVector;
 } // namespace llvm
 
 namespace mlir {
 
 namespace detail {
for (const auto &exprs : exprsList) {
     }
   }
 }
 
 inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
   map.print(os);
   return os;
 }
+
+// Return a bitvector where each bit set indicates a dimension that is not used
+// by any of the maps in the input array `maps`.
+llvm::SmallBitVector getUnusedDimsBitVector(ArrayRef<AffineMap> maps);
+
 } // namespace mlir
 
 namespace llvm {
 
 // AffineExpr hash just like pointers
 template <>
 struct DenseMapInfo<mlir::AffineMap> {
   static mlir::AffineMap getEmptyKey() {

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

if (!resVectorType || !accVectorType)
       return op.emitOpError("invalid accumulator/result vector shape");
 
     // Infer expected result vector type. Lhs + rhs map and lhs + rhs vector
     // types fully define the result vector type. This assumes the affine maps
     // are well-formed, which must have been verified already.
     MLIRContext *ctx = op.getContext();
     AffineMap lhsMap = op.getIndexingMaps()[0];
     AffineMap rhsMap = op.getIndexingMaps()[1];
+    if (getUnusedDimsBitVector({lhsMap, rhsMap}).any())
+      return op.emitOpError(
+          "expected all dimensions to be either a LHS or a RHS dimension");
     SmallVector<AffineExpr, 4> extents(lhsMap.getNumInputs());
     for (auto pair :
          {std::make_pair(lhsType, lhsMap), std::make_pair(rhsType, rhsMap)}) {
       VectorType v = pair.first;
       auto map = pair.second;
       for (unsigned idx = 0, e = v.getRank(); idx < e; ++idx) {
         unsigned pos = map.getDimPosition(idx);
         if (!extents[pos])
           extents[pos] = getAffineConstantExpr(v.getShape()[idx], ctx);
       }
     }
     if (!llvm::all_of(extents, [](AffineExpr e) { return e; }))
-      return op.emitOpError("expected all input dimensions to be used by "
-                            "either the LHS or the RHS");
+      return op.emitOpError("expected all dimensions to get an extent as "
+                            "either a LHS or a RHS dimension");
 
     AffineMap resMap = op.getIndexingMaps()[2];
     auto extentsMap = AffineMap::get(/*dimCount=*/extents.size(),
                                      /*symCount=*/0, extents, ctx);
     // Compose the resMap with the extentsMap, which is a constant map.
     AffineMap expectedMap = simplifyAffineMap(resMap.compose(extentsMap));
     assert(llvm::all_of(
                expectedMap.getResults(),

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

 #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"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "vector-to-vector"
 
 using namespace mlir;
 using namespace mlir::vector;
for (Value *operand : {&lhs, &rhs}) {
       changed = true;
     }
 
     if (!changed)
       return failure();
 
     // Determine which dims are usused, now that the maps have been composed
     // with the broadcast maps.
-    unsigned numDims = maps[0].getNumDims();
-    llvm::SmallBitVector unusedDims(numDims, true);
-    for (const auto &m : maps) {
-      for (unsigned i = 0; i < numDims; ++i) {
-        if (m.isFunctionOfDim(i))
-          unusedDims.reset(i);
-      }
-    }
+    llvm::SmallBitVector unusedDimsBitVector = getUnusedDimsBitVector(maps);
     // Compress unused dims.
     for (auto &m : maps)
-      m = compressDims(m, unusedDims);
+      m = compressDims(m, unusedDimsBitVector);
     // Compute the combined iterators.
     SmallVector<Attribute, 4> iterators;
-    for (unsigned i = 0; i < numDims; ++i) {
-      if (!unusedDims.test(i))
+    for (unsigned i = 0; i < unusedDimsBitVector.size(); ++i) {
+      if (!unusedDimsBitVector.test(i))
         iterators.push_back(contractOp.getIteratorTypes().getValue()[i]);
     }
     // Check that compressing unused dims isn't removing all reduction
     // iterators. For example, if the vector.contract had only one reduction
     // iterator and that was a unit-dimension created by a broadcast,
     // then we should bail here, otherwise we would create a contract without
     // a reduction iterator.
     if (!llvm::any_of(iterators, isReductionIterator))
       return failure();
-
+    // If the compressed maps have a dimension that is not used by either LHS or
+    // RHS then the ContractionOp verifier would fail.
+    if (getUnusedDimsBitVector({maps[0], maps[1]}).any())
+      return failure();
     rewriter.replaceOpWithNewOp<vector::ContractionOp>(
         contractOp, lhs, rhs, contractOp.getAcc(),
         rewriter.getAffineMapArrayAttr(maps), rewriter.getArrayAttr(iterators));
     return success();
   }
 };
 
 /// Reorders cast(broadcast) to broadcast(cast). This makes broadcast ops and

mlir/lib/IR/AffineMap.cpp

AffineMap mlir::compressDims(AffineMap map,
   SmallVector<AffineExpr> resultExprs;
   resultExprs.reserve(map.getNumResults());
   for (auto e : map.getResults())
     resultExprs.push_back(e.replaceDims(dimReplacements));
   return AffineMap::get(numDims, map.getNumSymbols(), resultExprs, context);
 }
 
 AffineMap mlir::compressUnusedDims(AffineMap map) {
-  llvm::SmallBitVector unusedDims(map.getNumDims(), true);
-  map.walkExprs([&](AffineExpr expr) {
-    if (auto dimExpr = expr.dyn_cast<AffineDimExpr>())
-      unusedDims.reset(dimExpr.getPosition());
-  });
-  return compressDims(map, unusedDims);
+  return compressDims(map, getUnusedDimsBitVector({map}));
 }
 
 static SmallVector<AffineMap>
 compressUnusedImpl(ArrayRef<AffineMap> maps,
                    llvm::function_ref<AffineMap(AffineMap)> compressionFun) {
   if (maps.empty())
     return SmallVector<AffineMap>();
   SmallVector<AffineExpr> allExprs;
return AffineMap::get(numDims, numSymbols, results,
                         maps.front().getContext());
 }
 
 AffineMap mlir::getProjectedMap(AffineMap map,
                                 const llvm::SmallBitVector &unusedDims) {
   return compressUnusedSymbols(compressDims(map, unusedDims));
 }
 
+llvm::SmallBitVector mlir::getUnusedDimsBitVector(ArrayRef<AffineMap> maps) {
+  unsigned numDims = maps[0].getNumDims();
+  llvm::SmallBitVector numDimsBitVector(numDims, true);
+  for (const auto &m : maps) {
+    for (unsigned i = 0; i < numDims; ++i) {
+      if (m.isFunctionOfDim(i))
+        numDimsBitVector.reset(i);
+    }
+  }
+  return numDimsBitVector;
+}
+
 //===----------------------------------------------------------------------===//
 // MutableAffineMap.
 //===----------------------------------------------------------------------===//
 
 MutableAffineMap::MutableAffineMap(AffineMap map)
     : results(map.getResults().begin(), map.getResults().end()),
       numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
       context(map.getContext()) {}

mlir/test/Dialect/Vector/invalid.mlir

// expected-error@+1 {{invalid accumulator/result vector shape, expected: 'vector<3x2xf32>'}}
   %0 = vector.contract #contraction_trait %arg0, %arg1, %arg2
     : vector<2x1xf32>, vector<1x3xf32> into vector<2x3xf32>
   return %0 : vector<2x3xf32>
 }
 
 // -----
 
 func.func @contract_with_dim_unused_by_lhs_and_rhs(%arg0 : vector<1x2xi32>, %arg1 : vector<2xi32>, %arg2 : vector<1xi32>) -> vector<1xi32> {
-// expected-error@+1 {{'vector.contract' op expected all input dimensions to be used by either the LHS or the RHS}}
+// expected-error@+1 {{'vector.contract' op expected all dimensions to be either a LHS or a RHS dimension}}
   %result = vector.contract {
     indexing_maps = [
       affine_map<(d0, d1, d2) -> (d0, d2)>,
       affine_map<(d0, d1, d2) -> (d2)>,
       affine_map<(d0, d1, d2) -> (d1)>
     ],
     iterator_types = ["reduction", "parallel", "reduction"],
     kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<1x2xi32>, vector<2xi32> into vector<1xi32>

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

 
 // Test that CombineContractBroadcast is not combining this case, as that would
 // result in dropping this contract's only reduction iterator.
 
 #map0 = affine_map<(d0, d1, d2) -> (d0, d1)>
 #map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
 #map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
 
+// CHECK-DAG: #[[$map0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[$map1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+
 // CHECK-LABEL: contract_broadcast_unit_dim_reduction_as_only_reduction
 //  CHECK-SAME: (%[[ARG0:.+]]: vector<8xi32>, %[[ARG1:.+]]: vector<8xi32>, %[[ARG2:.+]]: vector<8x8xi32>)
 //  CHECK: %[[BROADCAST0:.+]] = vector.broadcast %[[ARG0]] : vector<8xi32> to vector<1x8xi32>
 //  CHECK: %[[BROADCAST1:.+]] = vector.broadcast %[[ARG1]] : vector<8xi32> to vector<1x8xi32>
 //  CHECK: vector.contract
 //  CHECK-SAME: indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]]
 //  CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel"]
 //  CHECK-SAME: %[[BROADCAST0]], %[[BROADCAST1]], %[[ARG2]] : vector<1x8xi32>, vector<1x8xi32> into vector<8x8xi32>
 func.func @contract_broadcast_unit_dim_reduction_as_only_reduction(%arg0 : vector<8xi32>, %arg1 : vector<8xi32>, %arg2 : vector<8x8xi32>) -> vector<8x8xi32> {
     %0 = vector.broadcast %arg0 : vector<8xi32> to vector<1x8xi32>
     %1 = vector.broadcast %arg1 : vector<8xi32> to vector<1x8xi32>
     %result = vector.contract {
         indexing_maps = [#map0, #map1, #map2],
         iterator_types = ["reduction", "parallel", "parallel"],
         kind = #vector.kind<add>
       } %0, %1, %arg2 : vector<1x8xi32>, vector<1x8xi32> into vector<8x8xi32>
     return %result : vector<8x8xi32>
 }
 
+// -----
+
+// Test that CombineContractBroadcast is not combining this case, as that would
+// result in a dimension being unused in the LHS and RHS maps, which is illegal.
+
+#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1)>
+
+// CHECK-DAG: #[[$map0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$map1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d1)>
+
+// CHECK-LABEL: contract_broadcast_dimension_would_go_unused_in_lhs_rhs
+//  CHECK-SAME: (%[[ARG0:.+]]: vector<1x2xi32>, %[[ARG1:.+]]: vector<2xi32>, %[[ARG2:.+]]: vector<1xi32>)
+//  CHECK: %[[BROADCAST1:.+]] = vector.broadcast %[[ARG1]] : vector<2xi32> to vector<1x1x2xi32>
+//  CHECK: vector.contract
+//  CHECK-SAME: indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]]
+//  CHECK-SAME: iterator_types = ["reduction", "parallel", "reduction"]
+//  CHECK-SAME: %[[ARG0]], %[[BROADCAST1]], %[[ARG2]] : vector<1x2xi32>, vector<1x1x2xi32> into vector<1xi32>
+
+func.func @contract_broadcast_dimension_would_go_unused_in_lhs_rhs(%arg0 : vector<1x2xi32>, %arg1 : vector<2xi32>, %arg2 : vector<1xi32>) -> vector<1xi32> {
+  %1 = vector.broadcast %arg1 : vector<2xi32> to vector<1x1x2xi32>
+  %result = vector.contract {
+      indexing_maps = [#map0, #map1, #map2],
+      iterator_types = ["reduction", "parallel", "reduction"],
+      kind = #vector.kind<add>
+  } %arg0, %1, %arg2 : vector<1x2xi32>, vector<1x1x2xi32> into vector<1xi32>
+  return  %result : vector<1xi32>
+}
+
 //===----------------------------------------------------------------------===//
 // Reorder casting ops and vector ops. The casting ops have almost identical
 // pattern, so only arith.extsi op is tested.
 //===----------------------------------------------------------------------===//
 
 // -----
 
 func.func @broadcast_vector_extsi(%a : vector<4xi8>) -> vector<2x4xi32> {