Index: mlir/include/mlir/Dialect/Vector/VectorOps.td =================================================================== --- mlir/include/mlir/Dialect/Vector/VectorOps.td +++ mlir/include/mlir/Dialect/Vector/VectorOps.td @@ -461,40 +461,23 @@ 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 + with a size of `vector size / multiplicity`. Similarly to vector.tuple_get, this operation is used for progressive lowering and should be folded away before converting to LLVM. - - 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> + For instance: ``` - 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> - ``` - - 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 : 2] : vector<4xi32> to vector<2xi32> + // extracts values [1, 2] + %2 = vector.extract_map %0[%idx1 : 2] : vector<4xi32> to vector<2xi32> ``` Example: @@ -694,22 +677,36 @@ } 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, + I64Attr:$multiplicity)>, + 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 a size of `vector size * multiplicity`. Similarly to vector.tuple_get, this operation is used for progressive lowering and should be folded away before converting to LLVM. 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 : 2] : vector<4xi32> to vector<2xi32> + // extracts values [1, 2] + %2 = vector.extract_map %0[%idx1 : 2] : 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 : 2] : vector<2xi32> to vector<4xi32> + // insert [1, 2] into [x, x, x, x] and produce [x, 1, 2, x] + %4 = vector.insert_map %2, %0[%idx1 : 2] : vector<2xi32> to vector<4xi32> + ``` Example: ```mlir @@ -717,7 +714,7 @@ ``` }]; 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(); @@ -727,8 +724,8 @@ } }]; let assemblyFormat = [{ - $vector `,` $id `,` $multiplicity attr-dict `:` type($vector) `to` - type(results) + $vector `,` $dest `[` $id `:` $multiplicity `]` attr-dict + `:` type($vector) `to` type($result) }]; } Index: mlir/lib/Dialect/Vector/VectorOps.cpp =================================================================== --- mlir/lib/Dialect/Vector/VectorOps.cpp +++ mlir/lib/Dialect/Vector/VectorOps.cpp @@ -1191,11 +1191,13 @@ //===----------------------------------------------------------------------===// 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 resultType = VectorType::get(type.getNumElements() * multiplicity, type.getElementType()); - InsertMapOp::build(builder, result, resultType, vector, id, multiplicity); + InsertMapOp::build(builder, result, resultType, vector, dest, id, + multiplicity); } static LogicalResult verify(InsertMapOp op) { Index: mlir/lib/Dialect/Vector/VectorTransforms.cpp =================================================================== --- mlir/lib/Dialect/Vector/VectorTransforms.cpp +++ mlir/lib/Dialect/Vector/VectorTransforms.cpp @@ -2507,8 +2507,8 @@ DistributeOps ops; ops.extract = builder.create(loc, result, id, multiplicity); - ops.insert = - builder.create(loc, ops.extract, id, multiplicity); + ops.insert = builder.create(loc, ops.extract, result, id, + multiplicity); return ops; } @@ -2532,8 +2532,10 @@ Value newRead = vector_transfer_read(extract.getType(), read.memref(), indices, read.permutation_map(), read.padding(), ArrayAttr()); + Value dest = rewriter.create( + read.getLoc(), read.getType(), rewriter.getZeroAttr(read.getType())); newRead = rewriter.create( - read.getLoc(), newRead, extract.id(), extract.multiplicity()); + read.getLoc(), newRead, dest, extract.id(), extract.multiplicity()); rewriter.replaceOp(read, newRead); return success(); } Index: mlir/test/Dialect/Vector/vector-distribution.mlir =================================================================== --- mlir/test/Dialect/Vector/vector-distribution.mlir +++ mlir/test/Dialect/Vector/vector-distribution.mlir @@ -2,10 +2,11 @@ // CHECK-LABEL: func @distribute_vector_add // CHECK-SAME: (%[[ID:.*]]: index +// CHECK-NEXT: %[[ADDV:.*]] = addf %{{.*}}, %{{.*}} : vector<32xf32> // 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: %[[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]] : 32] : vector<1xf32> to 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> Index: mlir/test/lib/Transforms/TestVectorTransforms.cpp =================================================================== --- mlir/test/lib/Transforms/TestVectorTransforms.cpp +++ mlir/test/lib/Transforms/TestVectorTransforms.cpp @@ -146,7 +146,7 @@ Optional ops = distributPointwiseVectorOp( builder, op.getOperation(), func.getArgument(0), multiplicity); if (ops.hasValue()) { - SmallPtrSet extractOp({ops->extract}); + SmallPtrSet extractOp({ops->extract, ops->insert}); op.getResult().replaceAllUsesExcept(ops->insert.getResult(), extractOp); } });