[mlir][Vector] Disable 'vector.extract' folding for unsupported 0-D vectors

[dcaballe]

Summary

The vector.extract folding patterns do not support 0-D vectors
(actually, 0-D vector support couldn't even be implemented as a folding
pattern as it would require replacing vector.extract with a
vector.extractelement op). This patch is bailing out folding when 0-D
vectors are found.

[nicolasvasilache]

Note that ExtractOp does not support 0-D vectors at this time.
It seems to me that all these checks could be defensive asserts.

Separately, we may consider extending ExtractOp to support 0-D vectors, which would make sense given that ExtractElement returns elemental types, but we'll also need tests.

[dcaballe]

Good point! I can turn the check on ExtractOp into an assert. The key check, though, is the one on the defining op. That one may have 0-D vectors and we may end up creating a 0-D ExtracOp as a result of the folding. That's what the patch is aim to prevent.

[hanchung]

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

1628–1631

Can we move it after 0-D vectors check, then we are sure that the dest type is VectorType? We can still use cast and the cast has assertion semantics -- which is good to me.

In the context, the dest type must be a vector type because input operand type is not a 0-D vector.

[dcaballe]

Introduce asserts for ExtractOp

[nicolasvasilache]

ok, makes sense, can we have some tests here ?

[dcaballe]

Adding test. Some case can't be tested because other canonicalization patterns get in the way.

[nicolasvasilache]

thanks!

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

while (nextInsertOp || nextTransposeOp) {
     // Case 5: No intersection, we forward the extract to insertOp.dest().
     valueToExtractFrom = nextInsertOp.getDest();
     updateStateForNextIteration(valueToExtractFrom);
   }
   // If after all this we can fold, go for it.
   return tryToFoldExtractOpInPlace(valueToExtractFrom);
 }
 
+/// Returns true if the operation has a 0-D vector type operand or result.
+static bool hasZeroDimVectors(Operation *op) {
+  auto hasZeroDimVectorType = [](Type type) -> bool {
+    auto vecType = dyn_cast<VectorType>(type);
+    return vecType && vecType.getRank() == 0;
+  };
+
+  return llvm::any_of(op->getOperandTypes(), hasZeroDimVectorType) ||
+         llvm::any_of(op->getResultTypes(), hasZeroDimVectorType);
+}
+
 /// Fold extractOp with scalar result coming from BroadcastOp or SplatOp.
 static Value foldExtractFromBroadcast(ExtractOp extractOp) {
   Operation *defOp = extractOp.getVector().getDefiningOp();
   if (!defOp || !isa<vector::BroadcastOp, SplatOp>(defOp))
     return Value();
+
+  // 0-D vectors not supported.
+  assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
+  if (hasZeroDimVectors(defOp))
+    return Value();
+
   Value source = defOp->getOperand(0);
   if (extractOp.getType() == source.getType())
     return source;
   auto getRank = [](Type type) {
     return llvm::isa<VectorType>(type) ? llvm::cast<VectorType>(type).getRank()
                                        : 0;
   };
   // If splat or broadcast from a scalar, just return the source scalar.
static Value foldExtractFromBroadcast(ExtractOp extractOp) {
   return extractOp.getResult();
 }
 
 // Fold extractOp with source coming from ShapeCast op.
 static Value foldExtractFromShapeCast(ExtractOp extractOp) {
   auto shapeCastOp = extractOp.getVector().getDefiningOp<vector::ShapeCastOp>();
   if (!shapeCastOp)
     return Value();
+
+  // 0-D vectors not supported.
+  assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
+  if (hasZeroDimVectors(shapeCastOp))
+    return Value();
+
   // Get the nth dimension size starting from lowest dimension.
   auto getDimReverse = [](VectorType type, int64_t n) {
     return type.getShape().take_back(n + 1).front();
   };
   int64_t destinationRank =
       llvm::isa<VectorType>(extractOp.getType())
           ? llvm::cast<VectorType>(extractOp.getType()).getRank()
           : 0;

 }
 
 /// Fold an ExtractOp from ExtractStridedSliceOp.
 static Value foldExtractFromExtractStrided(ExtractOp extractOp) {
   auto extractStridedSliceOp =
       extractOp.getVector().getDefiningOp<vector::ExtractStridedSliceOp>();
   if (!extractStridedSliceOp)
     return Value();
+
+  // 0-D vectors not supported.
+  assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
+  if (hasZeroDimVectors(extractStridedSliceOp))
+    return Value();
+
   // Return if 'extractStridedSliceOp' has non-unit strides.
   if (extractStridedSliceOp.hasNonUnitStrides())
     return Value();
 
   // Trim offsets for dimensions fully extracted.
   auto sliceOffsets =
       extractVector<int64_t>(extractStridedSliceOp.getOffsets());
   while (!sliceOffsets.empty()) {
static Value foldExtractFromExtractStrided(ExtractOp extractOp) {
   // OpBuilder is only used as a helper to build an I64ArrayAttr.
   OpBuilder b(extractOp.getContext());
   extractOp->setAttr(ExtractOp::getPositionAttrStrName(),
                      b.getI64ArrayAttr(extractedPos));
   return extractOp.getResult();
 }
 
 /// Fold extract_op fed from a chain of insertStridedSlice ops.
-static Value foldExtractStridedOpFromInsertChain(ExtractOp op) {
-  int64_t destinationRank = llvm::isa<VectorType>(op.getType())
-                                ? llvm::cast<VectorType>(op.getType()).getRank()
-                                : 0;
-  auto insertOp = op.getVector().getDefiningOp<InsertStridedSliceOp>();
+static Value foldExtractStridedOpFromInsertChain(ExtractOp extractOp) {
+  int64_t destinationRank =
+      llvm::isa<VectorType>(extractOp.getType())
+          ? llvm::cast<VectorType>(extractOp.getType()).getRank()
+  auto insertOp = extractOp.getVector().getDefiningOp<InsertStridedSliceOp>();
+  if (!insertOp)
+    return Value();
+
+  // 0-D vectors not supported.
+  assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
+  if (hasZeroDimVectors(insertOp))
+    return Value();
+
   while (insertOp) {
     int64_t insertRankDiff = insertOp.getDestVectorType().getRank() -
                              insertOp.getSourceVectorType().getRank();
     if (destinationRank > insertOp.getSourceVectorType().getRank())
       return Value();
     auto insertOffsets = extractVector<int64_t>(insertOp.getOffsets());
-    auto extractOffsets = extractVector<int64_t>(op.getPosition());
+    auto extractOffsets = extractVector<int64_t>(extractOp.getPosition());
 
     if (llvm::any_of(insertOp.getStrides(), [](Attribute attr) {
           return llvm::cast<IntegerAttr>(attr).getInt() != 1;
         }))
       return Value();
     bool disjoint = false;
     SmallVector<int64_t, 4> offsetDiffs;
     for (unsigned dim = 0, e = extractOffsets.size(); dim < e; ++dim) {
if (!disjoint) {
       int64_t srcRankDiff =
           insertOp.getSourceVectorType().getRank() - destinationRank;
       for (int64_t i = 0; i < destinationRank; i++) {
         if (insertOp.getSourceVectorType().getDimSize(i + srcRankDiff) !=
             insertOp.getDestVectorType().getDimSize(i + srcRankDiff +
                                                     insertRankDiff))
           return Value();
       }
-      op.getVectorMutable().assign(insertOp.getSource());
+      extractOp.getVectorMutable().assign(insertOp.getSource());
       // OpBuilder is only used as a helper to build an I64ArrayAttr.
-      OpBuilder b(op.getContext());
-      op->setAttr(ExtractOp::getPositionAttrStrName(),
+      OpBuilder b(extractOp.getContext());
+      extractOp->setAttr(ExtractOp::getPositionAttrStrName(),
                          b.getI64ArrayAttr(offsetDiffs));
-      return op.getResult();
+      return extractOp.getResult();
     }
     // If the chunk extracted is disjoint from the chunk inserted, keep
     // looking in the insert chain.
     insertOp = insertOp.getDest().getDefiningOp<InsertStridedSliceOp>();
   }
   return Value();
 }
 

mlir/test/Dialect/Vector/canonicalize.mlir

 }
 
 // -----
 
 // CHECK-LABEL: fold_extract_shapecast_negative
 //       CHECK:   %[[V:.*]] = vector.shape_cast %{{.*}} : vector<16xf32> to vector<2x4x2xf32>
 //       CHECK:   %[[R:.*]] = vector.extract %[[V]][1] : vector<2x4x2xf32>
 //       CHECK:   return %[[R]] : vector<4x2xf32>
-func.func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>,
-                             %arg1 : vector<8x4x2xf32>) -> vector<4x2xf32> {
+func.func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>) -> vector<4x2xf32> {
   %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<2x4x2xf32>
   %r = vector.extract %0[1] : vector<2x4x2xf32>
   return %r : vector<4x2xf32>
 }
 
 // -----
 
+// CHECK-LABEL: dont_fold_0d_extract_shapecast
+//       CHECK:   %[[V:.*]] = vector.shape_cast %{{.*}} : vector<f32> to vector<1xf32>
+//       CHECK:   %[[R:.*]] = vector.extract %[[V]][0] : vector<1xf32>
+//       CHECK:   return %[[R]] : f32
+func.func @dont_fold_0d_extract_shapecast(%arg0 : vector<f32>) -> f32 {
+  %0 = vector.shape_cast %arg0 : vector<f32> to vector<1xf32>
+  %r = vector.extract %0[0] : vector<1xf32>
+  return %r : f32
+}
+
+// -----
+
 // CHECK-LABEL: dont_fold_expand_collapse
 //       CHECK:   %[[A:.*]] = vector.shape_cast %{{.*}} : vector<1x1x64xf32> to vector<1x1x8x8xf32>
 //       CHECK:   %[[B:.*]] = vector.shape_cast %{{.*}} : vector<1x1x8x8xf32> to vector<8x8xf32>
 //       CHECK:   return %[[B]] : vector<8x8xf32>
 func.func @dont_fold_expand_collapse(%arg0: vector<1x1x64xf32>) -> vector<8x8xf32> {
     %0 = vector.shape_cast %arg0 : vector<1x1x64xf32> to vector<1x1x8x8xf32>
     %1 = vector.shape_cast %0 : vector<1x1x8x8xf32> to vector<8x8xf32>
     return %1 : vector<8x8xf32>

 func.func @all_true_vector_mask(%a : vector<3x4xf32>) -> vector<3x4xf32> {
 //   CHECK-NOT:   vector.mask
 //       CHECK:   %[[ADD:.*]] = arith.addf %[[IN]], %[[IN]] : vector<3x4xf32>
 //       CHECK:   return %[[ADD]] : vector<3x4xf32>
   %all_true = vector.constant_mask [3, 4] : vector<3x4xi1>
   %0 = vector.mask %all_true { arith.addf %a, %a : vector<3x4xf32> } : vector<3x4xi1> -> vector<3x4xf32>
   return %0 : vector<3x4xf32>
 }
-