[mlir] Allow empty position in vector.insert and vector.extract

[springerm]

Summary

Such ops are no-ops and are folded to their respective source/vector operand.

Depends On D101745

[aartbik]

You rely on the folder to remove this, but would you mind adding some safety in the lowering code to LLVM (so that we don't get llvm.insertvalue %arg0, %arg1[] which will crash during lowering to LLVM IR)

mlir/lib/Dialect/Vector/VectorOps.cpp

1164

I think that form was used for the slightly longer folders. I would be perfectly okay inlining the two lines of code above in here. But if you like this style, I am okay with it too...

[springerm]

rebase

[springerm]

no change

[springerm]

no change

[nicolasvasilache]

You rely on the folder to remove this, but would you mind adding some safety in the lowering code to LLVM (so that we don't get llvm.insertvalue %arg0, %arg1[] which will crash during lowering to LLVM IR)

I'd avoid relying on the folder having to run beforehand.
We can just special case the empty position when lowering to LLVM and replace by the source during conversion to LLVM (rather than assert, which is how I interpret "adding some safety").

[aartbik]

In D101879#2746840, @nicolasvasilache wrote:

We can just special case the empty position when lowering to LLVM and replace by the source during conversion to LLVM (rather than assert, which is how I interpret "adding some safety").

Exactly. Either not rewriting or sipmly dropping to source (as the folder does) in the rewriting is exactly the kind of safety I had in mind.
Folders usually run, but I think it is not proper design to fully rely that that have run (especially when code can crash)

[springerm]

address comments

[springerm]

In D101879#2748591, @aartbik wrote:

In D101879#2746840, @nicolasvasilache wrote:

We can just special case the empty position when lowering to LLVM and replace by the source during conversion to LLVM (rather than assert, which is how I interpret "adding some safety").

Exactly. Either not rewriting or sipmly dropping to source (as the folder does) in the rewriting is exactly the kind of safety I had in mind.
Folders usually run, but I think it is not proper design to fully rely that that have run (especially when code can crash)

Done.

Btw, this commit was necessary for an earlier design of the ProgressiveVectorToSCF lowering. It's not needed in the current design anymore. I thought it may be useful to keep the commit anyway, to make the op more general, but if you think this doesn't really add anything, I can also abandon this commit.

[springerm]

address changes

[aartbik]

Let's keep it! I actually like having ops "generalize" to corner cases. Thanks for the change.

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

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

def Vector_InsertOp :
   ];
   let extraClassDeclaration = [{
     static StringRef getPositionAttrName() { return "position"; }
     Type getSourceType() { return source().getType(); }
     VectorType getDestVectorType() {
       return dest().getType().cast<VectorType>();
     }
   }];
+
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 }
 
 def Vector_InsertSlicesOp :
   Vector_Op<"insert_slices", [NoSideEffect]>,
     Arguments<(ins TupleOf<[AnyVector]>:$vectors, I64ArrayAttr:$sizes,
                    I64ArrayAttr:$strides)>,
     Results<(outs AnyVector)> {
   let summary = "vector insert slices operation";

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

matchAndRewrite(vector::ExtractOp extractOp, ArrayRef<Value> operands,
     auto resultType = extractOp.getResult().getType();
     auto llvmResultType = typeConverter->convertType(resultType);
     auto positionArrayAttr = extractOp.position();
 
     // Bail if result type cannot be lowered.
     if (!llvmResultType)
       return failure();
 
+    // Extract entire vector. Should be handled by folder, but just to be safe.
+    if (positionArrayAttr.empty()) {
+      rewriter.replaceOp(extractOp, adaptor.vector());
+      return success();
+    }
+
     // One-shot extraction of vector from array (only requires extractvalue).
     if (resultType.isa<VectorType>()) {
       Value extracted = rewriter.create<LLVM::ExtractValueOp>(
           loc, llvmResultType, adaptor.vector(), positionArrayAttr);
       rewriter.replaceOp(extractOp, extracted);
       return success();
     }
 
matchAndRewrite(vector::InsertOp insertOp, ArrayRef<Value> operands,
     auto destVectorType = insertOp.getDestVectorType();
     auto llvmResultType = typeConverter->convertType(destVectorType);
     auto positionArrayAttr = insertOp.position();
 
     // Bail if result type cannot be lowered.
     if (!llvmResultType)
       return failure();
 
+    // Overwrite entire vector with value. Should be handled by folder, but
+    // just to be safe.
+    if (positionArrayAttr.empty()) {
+      rewriter.replaceOp(insertOp, adaptor.source());
+      return success();
+    }
+
     // One-shot insertion of a vector into an array (only requires insertvalue).
     if (sourceType.isa<VectorType>()) {
       Value inserted = rewriter.create<LLVM::InsertValueOp>(
           loc, llvmResultType, adaptor.dest(), adaptor.source(),
           positionArrayAttr);
       rewriter.replaceOp(insertOp, inserted);
       return success();
     }

mlir/lib/Dialect/Vector/VectorOps.cpp

static ParseResult parseExtractOp(OpAsmParser &parser, OperationState &result) {
   Type resType = inferExtractOpResultType(vectorType, positionAttr);
   result.attributes = attrs;
   return failure(parser.resolveOperand(vector, type, result.operands) ||
                  parser.addTypeToList(resType, result.types));
 }
 
 static LogicalResult verify(vector::ExtractOp op) {
   auto positionAttr = op.position().getValue();
-  if (positionAttr.empty())
-    return op.emitOpError("expected non-empty position attribute");
   if (positionAttr.size() > static_cast<unsigned>(op.getVectorType().getRank()))
     return op.emitOpError(
         "expected position attribute of rank smaller than vector rank");
   for (auto en : llvm::enumerate(positionAttr)) {
     auto attr = en.value().dyn_cast<IntegerAttr>();
     if (!attr || attr.getInt() < 0 ||
         attr.getInt() >= op.getVectorType().getDimSize(en.index()))
       return op.emitOpError("expected position attribute #")
static Value foldExtractFromShapeCast(ExtractOp extractOp) {
   OpBuilder b(extractOp.getContext());
   extractOp->setAttr(ExtractOp::getPositionAttrName(),
                      b.getI64ArrayAttr(newPosition));
   extractOp.setOperand(shapeCastOp.source());
   return extractOp.getResult();
 }
 
 OpFoldResult ExtractOp::fold(ArrayRef<Attribute>) {
+  if (position().empty())
+    return vector();
   if (succeeded(foldExtractOpFromExtractChain(*this)))
     return getResult();
   if (succeeded(foldExtractOpFromTranspose(*this)))
     return getResult();
   if (auto val = foldExtractOpFromInsertChainAndTranspose(*this))
     return val;
   if (auto val = foldExtractFromBroadcast(*this))
     return val;
   if (auto val = foldExtractFromShapeCast(*this))
     return val;
   return OpFoldResult();
 }
 
 namespace {
 
 // If extractOp is only removing unit dimensions it can be transformed to a
 // shapecast.
 class ExtractToShapeCast final : public OpRewritePattern<ExtractOp> {
 public:
SmallVector<int64_t, 4> positionConstants =
       llvm::to_vector<4>(llvm::map_range(position, [](Value pos) {
         return pos.getDefiningOp<ConstantIndexOp>().getValue();
       }));
   build(builder, result, source, dest, positionConstants);
 }
 
 static LogicalResult verify(InsertOp op) {
   auto positionAttr = op.position().getValue();
-  if (positionAttr.empty())
-    return op.emitOpError("expected non-empty position attribute");
   auto destVectorType = op.getDestVectorType();
   if (positionAttr.size() > static_cast<unsigned>(destVectorType.getRank()))
     return op.emitOpError(
         "expected position attribute of rank smaller than dest vector rank");
   auto srcVectorType = op.getSourceType().dyn_cast<VectorType>();
   if (srcVectorType &&
       (static_cast<unsigned>(srcVectorType.getRank()) + positionAttr.size() !=
        static_cast<unsigned>(destVectorType.getRank())))

 
 } // namespace
 
 void InsertOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
   results.add<InsertToShapeCast>(context);
 }
 
+// Eliminates insert operations that produce values identical to their source
+// value. This happens when the source and destination vectors have identical
+// sizes.
+OpFoldResult vector::InsertOp::fold(ArrayRef<Attribute> operands) {
+  if (position().empty())
+    return source();
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // InsertSlicesOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(InsertSlicesOp op) {
   SmallVector<int64_t, 4> sizes;
   op.getSizes(sizes);
   SmallVector<int64_t, 4> strides;

mlir/test/Dialect/Vector/invalid.mlir

 
 func @extract_vector_type(%arg0: index) {
   // expected-error@+1 {{expected vector type}}
   %1 = vector.extract %arg0[] : index
 }
 
 // -----
 
-func @extract_position_empty(%arg0: vector<4x8x16xf32>) {
-  // expected-error@+1 {{expected non-empty position attribute}}
-  %1 = vector.extract %arg0[] : vector<4x8x16xf32>
-}
-
-// -----
-
 func @extract_position_rank_overflow(%arg0: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected position attribute of rank smaller than vector}}
   %1 = vector.extract %arg0[0, 0, 0, 0] : vector<4x8x16xf32>
 }
 
 // -----
 
 func @extract_position_rank_overflow_generic(%arg0: vector<4x8x16xf32>) {
func @insert_element_wrong_type(%arg0: i32, %arg1: vector<4xf32>) {
   %c = constant 3 : i32
   // expected-error@+1 {{'vector.insertelement' op failed to verify that source operand type matches element type of result}}
   %0 = "vector.insertelement" (%arg0, %arg1, %c) : (i32, vector<4xf32>, i32) -> (vector<4xf32>)
 }
 
 // -----
 
 func @insert_vector_type(%a: f32, %b: vector<4x8x16xf32>) {
-  // expected-error@+1 {{expected non-empty position attribute}}
-  %1 = vector.insert %a, %b[] : f32 into vector<4x8x16xf32>
-}
-
-// -----
-
-func @insert_vector_type(%a: f32, %b: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected position attribute of rank smaller than dest vector rank}}
   %1 = vector.insert %a, %b[3, 3, 3, 3, 3, 3] : f32 into vector<4x8x16xf32>
 }
 
 // -----
 
 func @insert_vector_type(%a: vector<4xf32>, %b: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected position attribute rank + source rank to match dest vector rank}}

mlir/test/Dialect/Vector/ops.mlir

func @extract_element(%a: vector<16xf32>) -> f32 {
   // CHECK:      %[[C15:.*]] = constant 15 : i32
   %c = constant 15 : i32
   // CHECK-NEXT: vector.extractelement %{{.*}}[%[[C15]] : i32] : vector<16xf32>
   %1 = vector.extractelement %a[%c : i32] : vector<16xf32>
   return %1 : f32
 }
 
 // CHECK-LABEL: @extract
-func @extract(%arg0: vector<4x8x16xf32>) -> (vector<8x16xf32>, vector<16xf32>, f32) {
+func @extract(%arg0: vector<4x8x16xf32>) -> (vector<4x8x16xf32>, vector<8x16xf32>, vector<16xf32>, f32) {
+  // CHECK: vector.extract {{.*}}[] : vector<4x8x16xf32>
+  %0 = vector.extract %arg0[] : vector<4x8x16xf32>
   // CHECK: vector.extract {{.*}}[3] : vector<4x8x16xf32>
   %1 = vector.extract %arg0[3] : vector<4x8x16xf32>
   // CHECK-NEXT: vector.extract {{.*}}[3, 3] : vector<4x8x16xf32>
   %2 = vector.extract %arg0[3, 3] : vector<4x8x16xf32>
   // CHECK-NEXT: vector.extract {{.*}}[3, 3, 3] : vector<4x8x16xf32>
   %3 = vector.extract %arg0[3, 3, 3] : vector<4x8x16xf32>
-  return %1, %2, %3 : vector<8x16xf32>, vector<16xf32>, f32
+  return %0, %1, %2, %3 : vector<4x8x16xf32>, vector<8x16xf32>, vector<16xf32>, f32
 }
 
 // CHECK-LABEL: @insert_element
 func @insert_element(%a: f32, %b: vector<16xf32>) -> vector<16xf32> {
   // CHECK:      %[[C15:.*]] = constant 15 : i32
   %c = constant 15 : i32
   // CHECK-NEXT: vector.insertelement %{{.*}}, %{{.*}}[%[[C15]] : i32] : vector<16xf32>
   %1 = vector.insertelement %a, %b[%c : i32] : vector<16xf32>
   return %1 : vector<16xf32>
 }
 
 // CHECK-LABEL: @insert
 func @insert(%a: f32, %b: vector<16xf32>, %c: vector<8x16xf32>, %res: vector<4x8x16xf32>) -> vector<4x8x16xf32> {
   // CHECK: vector.insert %{{.*}}, %{{.*}}[3] : vector<8x16xf32> into vector<4x8x16xf32>
   %1 = vector.insert %c, %res[3] : vector<8x16xf32> into vector<4x8x16xf32>
   // CHECK: vector.insert %{{.*}}, %{{.*}}[3, 3] : vector<16xf32> into vector<4x8x16xf32>
   %2 = vector.insert %b, %res[3, 3] : vector<16xf32> into vector<4x8x16xf32>
   // CHECK: vector.insert %{{.*}}, %{{.*}}[3, 3, 3] : f32 into vector<4x8x16xf32>
   %3 = vector.insert %a, %res[3, 3, 3] : f32 into vector<4x8x16xf32>
-  return %3 : vector<4x8x16xf32>
+  // CHECK: vector.insert %{{.*}}, %{{.*}}[] : vector<4x8x16xf32> into vector<4x8x16xf32>
+  %4 = vector.insert %3, %3[] : vector<4x8x16xf32> into vector<4x8x16xf32>
+  return %4 : vector<4x8x16xf32>
 }
 
 // CHECK-LABEL: @outerproduct
 func @outerproduct(%arg0: vector<4xf32>, %arg1: vector<8xf32>, %arg2: vector<4x8xf32>) -> vector<4x8xf32> {
   // CHECK: vector.outerproduct {{.*}} : vector<4xf32>, vector<8xf32>
   %0 = vector.outerproduct %arg0, %arg1 : vector<4xf32>, vector<8xf32>
   // CHECK: vector.outerproduct {{.*}}, {{.*}}, {{.*}} : vector<4xf32>, vector<8xf32>
   %1 = vector.outerproduct %arg0, %arg1, %arg2 : vector<4xf32>, vector<8xf32>