[mlir][Vector] Allow a 0-d for for vector transfer ops.

[nicolasvasilache]

Summary

This revision updates the op semantics, printer, parser and verifier to allow 0-d transfers.
Until 0-d vectors are available, such transfers have a special form that transits through vector<1xt>.
This is a stepping stone towards the longer term work of adding 0-d vectors and will help significantly reduce corner cases in vectorization.

Transformations and lowerings do not yet support this form, extensions will follow.

[dcaballe]

Just a quick comment for now... Before moving forward with this, would it make sense to send a small RFC to describe what a 0-d vector is, motivation, how it's different from a scalar or a vector<1xtype>. It seems like a significant addition to me and I think we have to make sure it aligns and composes well with tensors and memrefs.

[nicolasvasilache]

0-d vectors were discussed here: https://llvm.discourse.group/t/should-we-have-0-d-vectors/3097
There was quick consensus that the addition makes sense.

[springerm]

Note: vector::LoadOp and vector::StoreOp will also need to be updated. The LLVM lowering goes through these two ops. (There are also masked versions of those, but I think these can remain unchanged.)

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

1281

There could be some ambiguity with broadcasts, which are also encoded with zeros.

Maybe an affine map with no result would make more sense, since we actually have no dimension in the 0D case:

affine_map<()->()> or affine_map<()->(0)>

Not sure if this valid though...

mlir/lib/Dialect/Vector/VectorUtils.cpp

246–247

This affine map has no results.

[ThomasRaoux]

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

1281

+1, it sounds I believe affine_map<()->(0)> would be the most aligned with what we have now since the vector has a rank of 1 so we have 1 result and tensor has a rank of 0 so we should have 0 input dim.

[nicolasvasilache]

Address review.

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

1281

you're both right, thanks for catching this.

[nicolasvasilache]

In D111559#3056605, @springerm wrote:

Note: vector::LoadOp and vector::StoreOp will also need to be updated. The LLVM lowering goes through these two ops. (There are also masked versions of those, but I think these can remain unchanged.)

yes that will come in the followup

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

let description = [{
     %4 = vector.transfer_read %arg1[%c3, %c3], %vf0
       {permutation_map = (d0, d1)->(d0, d1)}
         : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
 
     // Read from a tensor with vector element type.
     %4 = vector.transfer_read %arg1[%c3, %c3], %vf0
       {permutation_map = (d0, d1)->(d0, d1)}
         : tensor<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
+
+    // Special encoding for 0-d transfer with 0-d tensor/memref, vector shape
+    // {1} and permutation_map () -> (0).
+    %0 = vector.transfer_read %arg0[], %f0 {permutation_map = affine_map<()->(0)>} :
+      tensor<f32>, vector<1xf32>
     ```
   }];
 
   let builders = [
     // Builder that sets padding to zero.
     OpBuilder<(ins "VectorType":$vector, "Value":$source,
       "ValueRange":$indices, "AffineMap":$permutationMap,
       CArg<"ArrayRef<bool>", "{}">:$inBounds)>,
let description = [{
     vector.transfer_write %4, %arg1[%c3, %c3]
       {permutation_map = (d0, d1)->(d0, d1)}
         : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
 
     // return a tensor where the vector is inserted into the source tensor.
     %5 = vector.transfer_write %4, %arg1[%c3, %c3]
       {permutation_map = (d0, d1)->(d0, d1)}
         : vector<1x1x4x3xf32>, tensor<?x?xvector<4x3xf32>>
+
+    // Special encoding for 0-d transfer with 0-d tensor/memref, vector shape
+    // {1} and permutation_map () -> (0).
+    %1 = vector.transfer_write %0, %arg0[] {permutation_map = affine_map<()->(0)>} :
+      vector<1xf32>, tensor<f32>
     ```
   }];
 
   let builders = [
     // Builder that sets permutation map to 'getMinorIdentityMap'.
     OpBuilder<(ins "Value":$vector, "Value":$source, "ValueRange":$indices,
       CArg<"ArrayRef<bool>", "{}">:$inBounds)>,
     OpBuilder<(ins "Value":$vector, "Value":$source, "ValueRange":$indices,

mlir/include/mlir/Interfaces/VectorInterfaces.td

InterfaceMethod<
       /*desc=*/"Return the permutation map.",
       /*retTy=*/"::mlir::AffineMap",
       /*methodName=*/"permutation_map",
       /*args=*/(ins),
       /*methodBody=*/"return $_op.permutation_map();"
       /*defaultImplementation=*/
     >,
     InterfaceMethod<
+      /*desc=*/[{
+        Returns true if op involves a 0-d tensor/memref and a vector
+        of shape {1}. This is temporary until we have 0-d vectors.
+        // TODO: turn this into 0-d vectors + empty permutation_map.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"isZeroD",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        if (getShapedType().getRank() > 0)
+          return false;
+        if (getVectorType().getShape() != ArrayRef<int64_t>{1})
+          return false;
+        AffineMap map = AffineMap::get(
+          /*numDims=*/0, /*numSymbols=*/0,
+          getAffineConstantExpr(0, $_op->getContext()));
+        if ($_op.permutation_map() != map)
+          return false;
+        return true;
+      }]
+    >,
+    InterfaceMethod<
       /*desc=*/[{ Returns true if the specified dimension is a broadcast. }],
       /*retTy=*/"bool",
       /*methodName=*/"isBroadcastDim",
       /*args=*/(ins "unsigned":$idx),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         auto expr = $_op.permutation_map().getResult(idx);
         return expr.template isa<::mlir::AffineConstantExpr>() &&
                expr.template dyn_cast<::mlir::AffineConstantExpr>().getValue() == 0;
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{ Returns true if at least one of the dimensions in the
                   permutation map is a broadcast.}],
       /*retTy=*/"bool",
       /*methodName=*/"hasBroadcastDim",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
+        // 0-d transfers are not considered broadcasts but they need to be 
+        // represented with a vector<1xt> until we have 0-d vectors.
+        if ($_op.isZeroD()) return false;
         for (unsigned i = 0; i < $_op.permutation_map().getNumResults(); ++i) {
           if ($_op.isBroadcastDim(i))
             return true;
         }
         return false;
       }]
     >,
     InterfaceMethod<

mlir/lib/Dialect/Vector/VectorOps.cpp

if (seen[dim.getPosition()]) {
           "requires a permutation_map that is a permutation (found one dim "
           "used more than once)");
     }
     seen[dim.getPosition()] = true;
   }
   return success();
 }
 
-static LogicalResult verifyTransferOp(Operation *op, ShapedType shapedType,
-                                      VectorType vectorType,
-                                      VectorType maskType,
-                                      AffineMap permutationMap,
-                                      ArrayAttr inBounds) {
+static LogicalResult
+verifyTransferOp(VectorTransferOpInterface op, ShapedType shapedType,
+                 VectorType vectorType, VectorType maskType,
+                 AffineMap permutationMap, ArrayAttr inBounds) {
+  if (shapedType.getRank() == 0 && !op.isZeroD())
+    return op->emitOpError("0-d transfer requires vector<1xt> shape and () -> "
+                           "(0) permutation_map");
+
   if (op->hasAttr("masked")) {
     return op->emitOpError("masked attribute has been removed. "
                            "Use in_bounds instead.");
   }
 
   if (!shapedType.isa<MemRefType, RankedTensorType>())
     return op->emitOpError(
         "requires source to be a memref or ranked tensor type");
if (auto vectorElementType = elementType.dyn_cast<VectorType>()) {
     if (maskType && expectedMaskType != maskType)
       return op->emitOpError("expects mask type consistent with permutation "
                              "map: ")
              << maskType;
   }
 
   if (permutationMap.getNumSymbols() != 0)
     return op->emitOpError("requires permutation_map without symbols");
-  if (permutationMap.getNumInputs() != shapedType.getRank())
+  // TODO: implement 0-d vector corner cases.
+  if (!op.isZeroD() && permutationMap.getNumInputs() != shapedType.getRank())
     return op->emitOpError("requires a permutation_map with input dims of the "
                            "same rank as the source type");
 
   if (inBounds) {
     if (permutationMap.getNumResults() != static_cast<int64_t>(inBounds.size()))
       return op->emitOpError("expects the optional in_bounds attr of same rank "
                              "as permutation_map results: ")
              << AffineMapAttr::get(permutationMap);
static LogicalResult verify(TransferReadOp op) {
   VectorType maskType = op.getMaskType();
   auto paddingType = op.padding().getType();
   auto permutationMap = op.permutation_map();
   auto sourceElementType = shapedType.getElementType();
 
   if (static_cast<int64_t>(op.indices().size()) != shapedType.getRank())
     return op.emitOpError("requires ") << shapedType.getRank() << " indices";
 
-  if (failed(verifyTransferOp(op.getOperation(), shapedType, vectorType,
-                              maskType, permutationMap,
-                              op.in_bounds() ? *op.in_bounds() : ArrayAttr())))
+  if (failed(
+          verifyTransferOp(cast<VectorTransferOpInterface>(op.getOperation()),
+                           shapedType, vectorType, maskType, permutationMap,
     return failure();
 
   if (auto sourceVectorElementType = sourceElementType.dyn_cast<VectorType>()) {
     // Source has vector element type.
     // Check that 'sourceVectorElementType' and 'paddingType' types match.
     if (sourceVectorElementType != paddingType)
       return op.emitOpError(
           "requires source element type and padding type to match.");
static bool isInBounds(TransferOp op, int64_t resultIdx, int64_t indicesIdx) {
   int64_t sourceSize = op.getShapedType().getDimSize(indicesIdx);
   int64_t vectorSize = op.getVectorType().getDimSize(resultIdx);
 
   return cstOp.getValue() + vectorSize <= sourceSize;
 }
 
 template <typename TransferOp>
 static LogicalResult foldTransferInBoundsAttribute(TransferOp op) {
+  // TODO: Be less conservative once we have 0-d vectors.
+  if (op.isZeroD())
+    return failure();
   AffineMap permutationMap = op.permutation_map();
   bool changed = false;
   SmallVector<bool, 4> newInBounds;
   newInBounds.reserve(op.getTransferRank());
   for (unsigned i = 0; i < op.getTransferRank(); ++i) {
     // Already marked as in-bounds, nothing to see here.
     if (op.isDimInBounds(i)) {
       newInBounds.push_back(true);
static LogicalResult verify(TransferWriteOp op) {
   if (llvm::size(op.indices()) != shapedType.getRank())
     return op.emitOpError("requires ") << shapedType.getRank() << " indices";
 
   // We do not allow broadcast dimensions on TransferWriteOps for the moment,
   // as the semantics is unclear. This can be revisited later if necessary.
   if (op.hasBroadcastDim())
     return op.emitOpError("should not have broadcast dimensions");
 
-  if (failed(verifyTransferOp(op.getOperation(), shapedType, vectorType,
-                              maskType, permutationMap,
-                              op.in_bounds() ? *op.in_bounds() : ArrayAttr())))
+  if (failed(
+          verifyTransferOp(cast<VectorTransferOpInterface>(op.getOperation()),
+                           shapedType, vectorType, maskType, permutationMap,
     return failure();
 
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
 /// Fold:
 /// ```

mlir/lib/Dialect/Vector/VectorUtils.cpp

 
 AffineMap mlir::getTransferMinorIdentityMap(ShapedType shapedType,
                                             VectorType vectorType) {
   int64_t elementVectorRank = 0;
   VectorType elementVectorType =
       shapedType.getElementType().dyn_cast<VectorType>();
   if (elementVectorType)
     elementVectorRank += elementVectorType.getRank();
+  // 0-d transfers are to/from tensor<t>/memref<t> and vector<1xt>.
+  // TODO: replace once we have 0-d vectors.
+  if (shapedType.getRank() == 0 &&
+      vectorType.getShape() == ArrayRef<int64_t>{1})
+    return AffineMap::get(
+        /*numDims=*/0, /*numSymbols=*/0,
+        getAffineConstantExpr(0, shapedType.getContext()));
   return AffineMap::getMinorIdentityMap(
       shapedType.getRank(), vectorType.getRank() - elementVectorRank,
       shapedType.getContext());
 }
 
 bool matcher::operatesOnSuperVectorsOf(Operation &op,
                                        VectorType subVectorType) {
   // First, extract the vector type and distinguish between:

mlir/test/Dialect/Vector/invalid.mlir

 }
 
 // -----
 
 func @insert_map_id(%v: vector<2x1xf32>, %v1: vector<4x32xf32>, %id : index) {
   // expected-error@+1 {{'vector.insert_map' op expected number of ids must match the number of dimensions distributed}}
   %0 = vector.insert_map %v, %v1[%id] : vector<2x1xf32> into vector<4x32xf32>
 }
+
+// -----
+
+func @vector_transfer_ops_0d(%arg0: tensor<f32>)
+  -> tensor<f32> {
+    %f0 = constant 0.0 : f32
+    // expected-error@+1 {{0-d transfer requires vector<1xt> shape and () -> (0) permutation_map}}
+    %0 = vector.transfer_read %arg0[], %f0 {permutation_map = affine_map<(d0)->(d0)>} :
+      tensor<f32>, vector<1xf32>
+    %1 = vector.transfer_write %0, %arg0[] {permutation_map = affine_map<()->(0)>} :
+      vector<1xf32>, tensor<f32>
+    return %1: tensor<f32>
+}

mlir/test/Dialect/Vector/ops.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 
+// CHECK-LABEL: func @vector_transfer_ops_0d(
+func @vector_transfer_ops_0d(%arg0: tensor<f32>, %arg1: memref<f32>)
+  -> tensor<f32> {
+    %f0 = constant 0.0 : f32
+    %0 = vector.transfer_read %arg0[], %f0 {permutation_map = affine_map<()->(0)>} :
+      tensor<f32>, vector<1xf32>
+    %1 = vector.transfer_write %0, %arg0[] {permutation_map = affine_map<()->(0)>} :
+      vector<1xf32>, tensor<f32>
+    %2 = vector.transfer_read %arg1[], %f0 {permutation_map = affine_map<()->(0)>} :
+      memref<f32>, vector<1xf32>
+    vector.transfer_write %2, %arg1[] {permutation_map = affine_map<()->(0)>} :
+      vector<1xf32>, memref<f32>
+    return %1: tensor<f32>
+}
+
 // CHECK-LABEL: func @vector_transfer_ops(
 func @vector_transfer_ops(%arg0: memref<?x?xf32>,
                           %arg1 : memref<?x?xvector<4x3xf32>>,
                           %arg2 : memref<?x?xvector<4x3xi32>>,
                           %arg3 : memref<?x?xvector<4x3xindex>>,
                           %arg4 : memref<?x?x?xf32>) {
   // CHECK: %[[C3:.*]] = constant 3 : index
   %c3 = constant 3 : index