fixup! [mlir][vector] Restrict use of 0-D vectors in vector.insert/vector.extract

Address Kunwar's comments

Author: banach-space

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

@@ -691,7 +691,7 @@ def Vector_ExtractOp :
      InferTypeOpAdaptorWithIsCompatible]> {
   let summary = "extract operation";
   let description = [{
-    Extracts an (n − k)-D subvector (the result) from an n-D vector at a
+    Extracts an (n − k)-D result sub-vector from an n-D source vector at a
     specified k-D position. When n = k, the result degenerates to a scalar
     element.
 
@@ -886,10 +886,9 @@ def Vector_InsertOp :
      AllTypesMatch<["dest", "result"]>]> {
   let summary = "insert operation";
   let description = [{
-    Inserts an n-D source vector (the value to store) into an (n + k)-D
-    destination vector at a specified k-D position. When n = 0, the source
-    degenerates to a scalar element inserted into the (0 + k)-D destination
-    vector.
+    Inserts an n-D value-to-store vector into an (n + k)-D destination vector
+    at a specified k-D position. When n = 0, value-to-store degenerates to
+    a scalar element inserted into the (0 + k)-D destination vector.
 
     Static and dynamic indices must be greater or equal to zero and less than
     the size of the corresponding dimension. The result is undefined if any

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

@@ -1294,10 +1294,6 @@ struct UnrollTransferReadConversion
 
   /// Rewrite the op: Unpack one dimension. Can handle masks, out-of-bounds
   /// accesses, and broadcasts and transposes in permutation maps.
-  ///
-  /// When unpacking rank-1 vectors (i.e. when the target rank is 0), replaces
-  /// `vector.transfer_read` with either `memref.load` or `tensor.extract` (for
-  /// MemRef and Tensor source, respectively).
   LogicalResult matchAndRewrite(TransferReadOp xferOp,
                                 PatternRewriter &rewriter) const override {
     if (xferOp.getVectorType().getRank() <= options.targetRank)
@@ -1345,32 +1341,20 @@ struct UnrollTransferReadConversion
 
             auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
 
-            // A value that's read after rank-reducing the original
-            // vector.transfer_read Op.
-            Value unpackedReadRes;
-            if (newXferVecType.getRank() != 0) {
-              // Unpacking Vector that's rank > 2
-              // (use vector.transfer_read to load a rank-reduced vector)
-              unpackedReadRes = b.create<vector::TransferReadOp>(
-                  loc, newXferVecType, xferOp.getBase(), xferIndices,
-                  AffineMapAttr::get(unpackedPermutationMap(b, xferOp)),
-                  xferOp.getPadding(), Value(), inBoundsAttr);
-              maybeAssignMask(b, xferOp,
-                              dyn_cast<vector::TransferReadOp>(
-                                  unpackedReadRes.getDefiningOp()),
-                              i);
-            } else {
-              // Unpacking Vector that's rank == 1
-              // (use memref.load/tensor.extract to load a scalar)
-              unpackedReadRes = dyn_cast<MemRefType>(xferOp.getBase().getType())
-                                    ? b.create<memref::LoadOp>(
-                                           loc, xferOp.getBase(), xferIndices)
-                                          .getResult()
-                                    : b.create<tensor::ExtractOp>(
-                                           loc, xferOp.getBase(), xferIndices)
-                                          .getResult();
+            auto newXferOp = b.create<vector::TransferReadOp>(
+                loc, newXferVecType, xferOp.getBase(), xferIndices,
+                AffineMapAttr::get(unpackedPermutationMap(b, xferOp)),
+                xferOp.getPadding(), Value(), inBoundsAttr);
+            maybeAssignMask(b, xferOp, newXferOp, i);
+
+            Value valToInser = newXferOp.getResult();
+            if (newXferVecType.getRank() == 0) {
+              // vector.insert does not accept rank-0 as the non-indexed
+              // argument. Extract the scalar before inserting.
+              valToInser = b.create<vector::ExtractOp>(loc, valToInser,
+                                                     SmallVector<int64_t>());
             }
-            return b.create<vector::InsertOp>(loc, unpackedReadRes, vec,
+            return b.create<vector::InsertOp>(loc, valToInser, vec,
                                               insertionIndices);
           },
           /*outOfBoundsCase=*/