[mlir][vector] Update doc strings for insert_map/extract_map and fix insert_map semantic

[ThomasRaoux]

Summary

Based on discourse discussion, fix the doc string and remove examples with wrong semantic. Also fix insert_map semantic by adding missing operand for vector we are inserting into.

[nicolasvasilache]

Thanks Thomas, this is a better separation of op semantics from transformation considerations indeed!

[nicolasvasilache]

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

463–464

Actually let's simplify this too:

Takes an 1-D vector and extract a sub-part of the vector starting at `id` with
a size of `vector size / multiplicity`.

No need to speculate on all the things that %id will represent in the op definition.
This makes it more in line with the LLVM boundary case: https://llvm.org/docs/LangRef.html#extractelement-instruction.

691–692

Same here please.

Thanks!

[aartbik]

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

463–464

nit: Takes ... and extracts (same verb form)

[nicolasvasilache]

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

717

Note this is missing the vector to insert into.
Syntax should also be symmetric to the extract counterpart.

%n = constant dense<x> : vector<4xi32>
%3 = vector.insert_map %1, %n[%idx0: 2] : vector<2xi32> to vector<4xi32>

[ThomasRaoux]

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

717

Makes sense, I think that was the key missing part.

[mehdi_amini]

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

717

Ah that changes quite the picture!

But why do we need a multiplicity attribute now? We have the type of the vector we insert into: don't we need only the index?

[ThomasRaoux]

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

717

Good point, this is redundant right now. I think there was some ideas behind it to be able to handle non aligned cases but this is definitely premature at this point. Removing it to make things clearer for now.

[ThomasRaoux]

Remove multiplicity attribute for now as it is redundant.

[mehdi_amini]

Could we also expand the doc for these op to explain how they differ from vector.insert and vector.extract?

[ThomasRaoux]

Add docs on difference with ExtractOp/InsertOp

[ThomasRaoux]

In D89563#2348962, @mehdi_amini wrote:

Could we also expand the doc for these op to explain how they differ from vector.insert and vector.extract?

Done.

[mehdi_amini]

Thanks!

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

469

471

700

(same)

[nicolasvasilache]

Thanks Thomas!

[nicolasvasilache]

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

719

Sorry I mistyped, can we s/to/into to be more consistent with other vector ops?

[ThomasRaoux]

Address more review comments

[ThomasRaoux]

Thanks for the review

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

719

Done.

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

def Vector_ExtractSlicesOp :
   let assemblyFormat = [{
     $vector `,` $sizes `,` $strides attr-dict `:` type($vector) `into`
     type(results)
   }];
 }
 
 def Vector_ExtractMapOp :
   Vector_Op<"extract_map", [NoSideEffect]>,
-    Arguments<(ins AnyVector:$vector, Index:$id, I64Attr:$multiplicity)>,
+    Arguments<(ins AnyVector:$vector, Index:$id)>,
     Results<(outs AnyVector)> {
   let summary = "vector extract map operation";
   let description = [{
-    Takes an 1-D vector and extract a sub-part of the vector starting at id with
-    a size of `vector size / multiplicity`. This maps a given multiplicity of
-    the vector to a Value such as a loop induction variable or an SPMD id.
+    Takes an 1-D vector and extracts a sub-part of the vector starting at id.
 
     Similarly to vector.tuple_get, this operation is used for progressive
     lowering and should be folded away before converting to LLVM.
 
+    It is different than `vector.extract_slice` and
+    `vector.extract_strided_slice` as it takes a Value as index instead of an
+    attribute. Also in the future it is meant to support extracting along any
+    dimensions and not only the most major ones.
 
-    For instance, the following code:
-    ```mlir
-    %a = vector.transfer_read %A[%c0]: memref<32xf32>, vector<32xf32>
-    %b = vector.transfer_read %B[%c0]: memref<32xf32>, vector<32xf32>
-    %c = addf %a, %b: vector<32xf32>
-    vector.transfer_write %c, %C[%c0]: memref<32xf32>, vector<32xf32>
-    ```
-    can be rewritten to:
-    ```mlir
-    %a = vector.transfer_read %A[%c0]: memref<32xf32>, vector<32xf32>
-    %b = vector.transfer_read %B[%c0]: memref<32xf32>, vector<32xf32>
-    %ea = vector.extract_map %a[%id : 32] : vector<32xf32> to vector<1xf32>
-    %eb = vector.extract_map %b[%id : 32] : vector<32xf32> to vector<1xf32>
-    %ec = addf %ea, %eb : vector<1xf32>
-    %c = vector.insert_map %ec, %id, 32 : vector<1xf32> to vector<32xf32>
-    vector.transfer_write %c, %C[%c0]: memref<32xf32>, vector<32xf32>
+    For instance:
     ```
-
-    Where %id can be an induction variable or an SPMD id going from 0 to 31.
-
-    And then be rewritten to:
-    ```mlir
-    %a = vector.transfer_read %A[%id]: memref<32xf32>, vector<1xf32>
-    %b = vector.transfer_read %B[%id]: memref<32xf32>, vector<1xf32>
-    %c = addf %a, %b: vector<1xf32>
-    vector.transfer_write %c, %C[%id]: memref<32xf32>, vector<1xf32>
+    // dynamic computation producing the value 0 of index type
+    %idx0 = ... : index
+    // dynamic computation producing the value 1 of index type
+    %idx1 = ... : index
+    %0 = constant dense<0, 1, 2, 3>: vector<4xi32>
+    // extracts values [0, 1]
+    %1 = vector.extract_map %0[%idx0] : vector<4xi32> to vector<2xi32>
+    // extracts values [1, 2]
+    %2 = vector.extract_map %0[%idx1] : vector<4xi32> to vector<2xi32>
     ```
 
     Example:
 
     ```mlir
-    %ev = vector.extract_map %v[%id:32] : vector<32xf32> to vector<1xf32>
+    %ev = vector.extract_map %v[%id] : vector<32xf32> to vector<1xf32>
     ```
   }];
   let builders = [OpBuilder<
     "Value vector, Value id, int64_t multiplicity">];
   let extraClassDeclaration = [{
     VectorType getSourceVectorType() {
       return vector().getType().cast<VectorType>();
     }
     VectorType getResultType() {
       return getResult().getType().cast<VectorType>();
     }
+    int64_t multiplicity() {
+      return getSourceVectorType().getNumElements() /
+        getResultType().getNumElements();
+    }
   }];
   let assemblyFormat = [{
-    $vector `[` $id `:` $multiplicity `]` attr-dict `:` type($vector) `to`
-    type(results)
+    $vector `[` $id `]` attr-dict `:` type($vector) `to` type(results)
   }];
 
   let hasFolder = 1;
 }
 
 def Vector_FMAOp :
   Op<Vector_Dialect, "fma", [NoSideEffect,
                              AllTypesMatch<["lhs", "rhs", "acc", "result"]>]>,
def Vector_InsertSlicesOp :
   }];
   let assemblyFormat = [{
     $vectors `,` $sizes `,` $strides attr-dict `:` type($vectors) `into`
     type(results)
   }];
 }
 
 def Vector_InsertMapOp :
-  Vector_Op<"insert_map", [NoSideEffect]>,
-    Arguments<(ins AnyVector:$vector, Index:$id, I64Attr:$multiplicity)>,
-    Results<(outs AnyVector)> {
+  Vector_Op<"insert_map", [NoSideEffect, AllTypesMatch<["dest", "result"]>]>,
+    Arguments<(ins AnyVector:$vector, AnyVector:$dest, Index:$id)>,
+    Results<(outs AnyVector:$result)> {
   let summary = "vector insert map operation";
   let description = [{
-    insert an 1-D vector and within a larger vector starting at id. The new
-    vector created will have a size of `vector size * multiplicity`. This
-    represents how a sub-part of the vector is written for a given Value such as
-    a loop induction variable or an SPMD id.
+    Inserts a 1-D vector and within a larger vector starting at id. The new
+    vector created will have the same size as the destination operand vector.
 
     Similarly to vector.tuple_get, this operation is used for progressive
     lowering and should be folded away before converting to LLVM.
 
+    It is different than `vector.insert` and `vector.insert_strided_slice` as it
+    takes a Value as index instead of an attribute. Also in the future it is
+    meant to support inserting along any dimensions and not only the most major
+    ones.
+
     This operations is meant to be used in combination with vector.extract_map.
-    See example in extract.map description.
 
+    For instance:
+    ```
+    // dynamic computation producing the value 0 of index type
+    %idx0 = ... : index
+    // dynamic computation producing the value 1 of index type
+    %idx1 = ... : index /
+    %0 = constant dense<0, 1, 2, 3>: vector<4xi32>
+    // extracts values [0, 1]
+    %1 = vector.extract_map %0[%idx0] : vector<4xi32> to vector<2xi32>
+    // extracts values [1, 2]
+    %2 = vector.extract_map %0[%idx1] : vector<4xi32> to vector<2xi32>
+    // insert [0, 1] into [x, x, x, x] and produce [0, 1, x, x]
+    %3 = vector.insert_map %1, %0[%idx0] : vector<2xi32> into vector<4xi32>
+    // insert [1, 2] into [x, x, x, x] and produce [x, 1, 2, x]
+    %4 = vector.insert_map %2, %0[%idx1] : vector<2xi32> into vector<4xi32>
+    ```
     Example:
 
     ```mlir
-    %v = vector.insert_map %ev, %id, 32 : vector<1xf32> to vector<32xf32>
+    %v = vector.insert_map %ev %v[%id] : vector<1xf32> into vector<32xf32>
     ```
   }];
   let builders = [OpBuilder<
-    "Value vector, Value id, int64_t multiplicity">];
+    "Value vector, Value dest, Value id, int64_t multiplicity">];
   let extraClassDeclaration = [{
     VectorType getSourceVectorType() {
       return vector().getType().cast<VectorType>();
     }
     VectorType getResultType() {
       return getResult().getType().cast<VectorType>();
     }
+    int64_t multiplicity() {
+      return getResultType().getNumElements() /
+        getSourceVectorType().getNumElements();
+    }
   }];
   let assemblyFormat = [{
-    $vector `,` $id `,` $multiplicity attr-dict `:` type($vector) `to`
-    type(results)
+    $vector `,` $dest `[` $id `]` attr-dict
+      `:` type($vector) `into` type($result)
   }];
 }
 
 def Vector_InsertStridedSliceOp :
   Vector_Op<"insert_strided_slice", [NoSideEffect,
     PredOpTrait<"operand #0 and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
     AllTypesMatch<["dest", "res"]>]>,

mlir/lib/Dialect/Vector/VectorOps.cpp

 // ExtractMapOp
 //===----------------------------------------------------------------------===//
 
 void ExtractMapOp::build(OpBuilder &builder, OperationState &result,
                          Value vector, Value id, int64_t multiplicity) {
   VectorType type = vector.getType().cast<VectorType>();
   VectorType resultType = VectorType::get(type.getNumElements() / multiplicity,
                                           type.getElementType());
-  ExtractMapOp::build(builder, result, resultType, vector, id, multiplicity);
+  ExtractMapOp::build(builder, result, resultType, vector, id);
 }
 
 static LogicalResult verify(ExtractMapOp op) {
   if (op.getSourceVectorType().getShape().size() != 1 ||
       op.getResultType().getShape().size() != 1)
     return op.emitOpError("expects source and destination vectors of rank 1");
-  if (op.getResultType().getNumElements() * (int64_t)op.multiplicity() !=
-      op.getSourceVectorType().getNumElements())
-    return op.emitOpError("vector sizes mismatch. Source size must be equal "
-                          "to destination size * multiplicity");
+  if (op.getSourceVectorType().getNumElements() %
+          op.getResultType().getNumElements() !=
+      0)
+    return op.emitOpError(
+        "source vector size must be a multiple of destination vector size");
   return success();
 }
 
 OpFoldResult ExtractMapOp::fold(ArrayRef<Attribute> operands) {
   auto insert = vector().getDefiningOp<vector::InsertMapOp>();
   if (insert == nullptr || multiplicity() != insert.multiplicity() ||
       id() != insert.id())
     return {};
void InsertSlicesOp::getStrides(SmallVectorImpl<int64_t> &results) {
   populateFromInt64AttrArray(strides(), results);
 }
 
 //===----------------------------------------------------------------------===//
 // InsertMapOp
 //===----------------------------------------------------------------------===//
 
 void InsertMapOp::build(OpBuilder &builder, OperationState &result,
-                        Value vector, Value id, int64_t multiplicity) {
+                        Value vector, Value dest, Value id,
+                        int64_t multiplicity) {
   VectorType type = vector.getType().cast<VectorType>();
   VectorType resultType = VectorType::get(type.getNumElements() * multiplicity,
                                           type.getElementType());
-  InsertMapOp::build(builder, result, resultType, vector, id, multiplicity);
+  InsertMapOp::build(builder, result, resultType, vector, dest, id);
 }
 
 static LogicalResult verify(InsertMapOp op) {
   if (op.getSourceVectorType().getShape().size() != 1 ||
       op.getResultType().getShape().size() != 1)
     return op.emitOpError("expected source and destination vectors of rank 1");
-  if ((int64_t)op.multiplicity() * op.getSourceVectorType().getNumElements() !=
-      op.getResultType().getNumElements())
+  if (op.getResultType().getNumElements() %
+          op.getSourceVectorType().getNumElements() !=
+      0)
     return op.emitOpError(
-        "vector sizes mismatch. Destination size must be equal "
-        "to source size * multiplicity");
+        "destination vector size must be a multiple of source vector size");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // InsertStridedSliceOp
 //===----------------------------------------------------------------------===//
 
 void InsertStridedSliceOp::build(OpBuilder &builder, OperationState &result,

mlir/lib/Dialect/Vector/VectorTransforms.cpp

mlir::vector::distributPointwiseVectorOp(OpBuilder &builder, Operation *op,
   VectorType type = op->getResult(0).getType().dyn_cast<VectorType>();
   // Currently only support distributing 1-D vectors of size multiple of the
   // given multiplicty. To handle more sizes we would need to support masking.
   if (!type || type.getRank() != 1 || type.getNumElements() % multiplicity != 0)
     return {};
   DistributeOps ops;
   ops.extract =
       builder.create<vector::ExtractMapOp>(loc, result, id, multiplicity);
-  ops.insert =
-      builder.create<vector::InsertMapOp>(loc, ops.extract, id, multiplicity);
+  ops.insert = builder.create<vector::InsertMapOp>(loc, ops.extract, result, id,
+                                                   multiplicity);
   return ops;
 }
 
 struct TransferReadExtractPattern
     : public OpRewritePattern<vector::TransferReadOp> {
   TransferReadExtractPattern(MLIRContext *context)
       : OpRewritePattern<vector::TransferReadOp>(context) {}
   LogicalResult matchAndRewrite(vector::TransferReadOp read,
                                 PatternRewriter &rewriter) const override {
     if (!read.getResult().hasOneUse())
       return failure();
     auto extract =
         dyn_cast<vector::ExtractMapOp>(*read.getResult().getUsers().begin());
     if (!extract)
       return failure();
     edsc::ScopedContext scope(rewriter, read.getLoc());
     using mlir::edsc::op::operator+;
     using namespace mlir::edsc::intrinsics;
     SmallVector<Value, 4> indices(read.indices().begin(), read.indices().end());
     indices.back() = indices.back() + extract.id();
     Value newRead = vector_transfer_read(extract.getType(), read.memref(),
                                          indices, read.permutation_map(),
                                          read.padding(), ArrayAttr());
+    Value dest = rewriter.create<ConstantOp>(
+        read.getLoc(), read.getType(), rewriter.getZeroAttr(read.getType()));
     newRead = rewriter.create<vector::InsertMapOp>(
-        read.getLoc(), newRead, extract.id(), extract.multiplicity());
+        read.getLoc(), newRead, dest, extract.id(), extract.multiplicity());
     rewriter.replaceOp(read, newRead);
     return success();
   }
 };
 
 struct TransferWriteInsertPattern
     : public OpRewritePattern<vector::TransferWriteOp> {
   TransferWriteInsertPattern(MLIRContext *context)

mlir/test/Dialect/Vector/invalid.mlir

func @compress_dim_mask_mismatch(%base: memref<?xf32>, %mask: vector<17xi1>, %value: vector<16xf32>) {
   // expected-error@+1 {{'vector.compressstore' op expected value dim to match mask dim}}
   vector.compressstore %base, %mask, %value : memref<?xf32>, vector<17xi1>, vector<16xf32>
 }
 
 // -----
 
 func @extract_map_rank(%v: vector<2x32xf32>, %id : index) {
   // expected-error@+1 {{'vector.extract_map' op expects source and destination vectors of rank 1}}
-  %0 = vector.extract_map %v[%id : 32] : vector<2x32xf32> to vector<2x1xf32>
+  %0 = vector.extract_map %v[%id] : vector<2x32xf32> to vector<2x1xf32>
 }
 
 // -----
 
 func @extract_map_size(%v: vector<63xf32>, %id : index) {
-  // expected-error@+1 {{'vector.extract_map' op vector sizes mismatch. Source size must be equal to destination size * multiplicity}}
-  %0 = vector.extract_map %v[%id : 32] : vector<63xf32> to vector<2xf32>
+  // expected-error@+1 {{'vector.extract_map' op source vector size must be a multiple of destination vector size}}
+  %0 = vector.extract_map %v[%id] : vector<63xf32> to vector<2xf32>
 }
 
 // -----
 
-func @insert_map_rank(%v: vector<2x1xf32>, %id : index) {
+func @insert_map_rank(%v: vector<2x1xf32>, %v1: vector<2x32xf32>, %id : index) {
   // expected-error@+1 {{'vector.insert_map' op expected source and destination vectors of rank 1}}
-  %0 = vector.insert_map %v, %id, 32 : vector<2x1xf32> to vector<2x32xf32>
+  %0 = vector.insert_map %v, %v1[%id] : vector<2x1xf32> into vector<2x32xf32>
 }
 
 // -----
 
-func @insert_map_size(%v: vector<1xf32>, %id : index) {
-  // expected-error@+1 {{'vector.insert_map' op vector sizes mismatch. Destination size must be equal to source size * multiplicity}}
-  %0 = vector.insert_map %v, %id, 32 : vector<1xf32> to vector<64xf32>
+func @insert_map_size(%v: vector<3xf32>, %v1: vector<64xf32>, %id : index) {
+  // expected-error@+1 {{'vector.insert_map' op destination vector size must be a multiple of source vector size}}
+  %0 = vector.insert_map %v, %v1[%id] : vector<3xf32> into vector<64xf32>
 }

mlir/test/Dialect/Vector/ops.mlir

func @expand_and_compress(%base: memref<?xf32>, %mask: vector<16xi1>, %passthru: vector<16xf32>) {
   %0 = vector.expandload %base, %mask, %passthru : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
   // CHECK: vector.compressstore %{{.*}}, %{{.*}}, %[[X]] : memref<?xf32>, vector<16xi1>, vector<16xf32>
   vector.compressstore %base, %mask, %0 : memref<?xf32>, vector<16xi1>, vector<16xf32>
   return
 }
 
 // CHECK-LABEL: @extract_insert_map
 func @extract_insert_map(%v: vector<32xf32>, %id : index) -> vector<32xf32> {
-  // CHECK: %[[V:.*]] = vector.extract_map %{{.*}}[%{{.*}} : 16] : vector<32xf32> to vector<2xf32>
-  %vd = vector.extract_map %v[%id : 16] : vector<32xf32> to vector<2xf32>
-  // CHECK: %[[R:.*]] = vector.insert_map %[[V]], %{{.*}}, 16 : vector<2xf32> to vector<32xf32>
-  %r = vector.insert_map %vd, %id, 16 : vector<2xf32> to vector<32xf32>
+  // CHECK: %[[V:.*]] = vector.extract_map %{{.*}}[%{{.*}}] : vector<32xf32> to vector<2xf32>
+  %vd = vector.extract_map %v[%id] : vector<32xf32> to vector<2xf32>
+  // CHECK: %[[R:.*]] = vector.insert_map %[[V]], %{{.*}}[%{{.*}}] : vector<2xf32> into vector<32xf32>
+  %r = vector.insert_map %vd, %v[%id] : vector<2xf32> into vector<32xf32>
   // CHECK: return %[[R]] : vector<32xf32>
   return %r : vector<32xf32>
 }
 

mlir/test/Dialect/Vector/vector-distribution.mlir

 // RUN: mlir-opt %s -test-vector-distribute-patterns=distribution-multiplicity=32 | FileCheck %s
 
 // CHECK-LABEL: func @distribute_vector_add
 //  CHECK-SAME: (%[[ID:.*]]: index
-//  CHECK-NEXT:    %[[EXA:.*]] = vector.extract_map %{{.*}}[%[[ID]] : 32] : vector<32xf32> to vector<1xf32>
-//  CHECK-NEXT:    %[[EXB:.*]] = vector.extract_map %{{.*}}[%[[ID]] : 32] : vector<32xf32> to vector<1xf32>
+//  CHECK-NEXT:    %[[ADDV:.*]] = addf %{{.*}}, %{{.*}} : vector<32xf32>
+//  CHECK-NEXT:    %[[EXA:.*]] = vector.extract_map %{{.*}}[%[[ID]]] : vector<32xf32> to vector<1xf32>
+//  CHECK-NEXT:    %[[EXB:.*]] = vector.extract_map %{{.*}}[%[[ID]]] : vector<32xf32> to vector<1xf32>
 //  CHECK-NEXT:    %[[ADD:.*]] = addf %[[EXA]], %[[EXB]] : vector<1xf32>
-//  CHECK-NEXT:    %[[INS:.*]] = vector.insert_map %[[ADD]], %[[ID]], 32 : vector<1xf32> to vector<32xf32>
+//  CHECK-NEXT:    %[[INS:.*]] = vector.insert_map %[[ADD]], %[[ADDV]][%[[ID]]] : vector<1xf32> into vector<32xf32>
 //  CHECK-NEXT:    return %[[INS]] : vector<32xf32>
 func @distribute_vector_add(%id : index, %A: vector<32xf32>, %B: vector<32xf32>) -> vector<32xf32> {
   %0 = addf %A, %B : vector<32xf32>
   return %0: vector<32xf32>
 }
 
 // CHECK-LABEL: func @vector_add_read_write
 //  CHECK-SAME: (%[[ID:.*]]: index

mlir/test/lib/Transforms/TestVectorTransforms.cpp

void runOnFunction() override {
     MLIRContext *ctx = &getContext();
     OwningRewritePatternList patterns;
     FuncOp func = getFunction();
     func.walk([&](AddFOp op) {
       OpBuilder builder(op);
       Optional<mlir::vector::DistributeOps> ops = distributPointwiseVectorOp(
           builder, op.getOperation(), func.getArgument(0), multiplicity);
       if (ops.hasValue()) {
-        SmallPtrSet<Operation *, 1> extractOp({ops->extract});
+        SmallPtrSet<Operation *, 1> extractOp({ops->extract, ops->insert});
         op.getResult().replaceAllUsesExcept(ops->insert.getResult(), extractOp);
       }
     });
     patterns.insert<PointwiseExtractPattern>(ctx);
     populateVectorToVectorTransformationPatterns(patterns, ctx);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };