[mlir][vector] Use DenseI64ArrayAttr for ExtractOp/InsertOp positions

`DenseI64ArrayAttr` provides a better API than `I64ArrayAttr`. E.g., accessors returning `ArrayRef<int64_t>` (instead of `ArrayAttr`) are generated.

Differential Revision: https://reviews.llvm.org/D156684

Author: matthias-springer

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

@@ -573,7 +573,7 @@ def Vector_ExtractOp :
      PredOpTrait<"operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      InferTypeOpAdaptorWithIsCompatible]>,
-    Arguments<(ins AnyVectorOfAnyRank:$vector, I64ArrayAttr:$position)>,
+    Arguments<(ins AnyVectorOfAnyRank:$vector, DenseI64ArrayAttr:$position)>,
     Results<(outs AnyType)> {
   let summary = "extract operation";
   let description = [{
@@ -589,7 +589,6 @@ def Vector_ExtractOp :
     ```
   }];
   let builders = [
-    OpBuilder<(ins "Value":$source, "ArrayRef<int64_t>":$position)>,
     // Convenience builder which assumes the values in `position` are defined by
     // ConstantIndexOp.
     OpBuilder<(ins "Value":$source, "ValueRange":$position)>
@@ -689,7 +688,7 @@ def Vector_InsertOp :
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      AllTypesMatch<["dest", "res"]>]>,
-     Arguments<(ins AnyType:$source, AnyVectorOfAnyRank:$dest, I64ArrayAttr:$position)>,
+     Arguments<(ins AnyType:$source, AnyVectorOfAnyRank:$dest, DenseI64ArrayAttr:$position)>,
      Results<(outs AnyVectorOfAnyRank:$res)> {
   let summary = "insert operation";
   let description = [{
@@ -711,8 +710,6 @@ def Vector_InsertOp :
   }];
 
   let builders = [
-    OpBuilder<(ins "Value":$source, "Value":$dest,
-      "ArrayRef<int64_t>":$position)>,
     // Convenience builder which assumes all values are constant indices.
     OpBuilder<(ins "Value":$source, "Value":$dest, "ValueRange":$position)>
   ];

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

@@ -807,8 +807,7 @@ createNonLdMatrixLoads(RewriterBase &rewriter, vector::TransferReadOp op,
 
       Value el = rewriter.create<vector::LoadOp>(loc, loadedElType,
                                                  op.getSource(), newIndices);
-      result = rewriter.create<vector::InsertOp>(loc, el, result,
-                                                 rewriter.getI64ArrayAttr(i));
+      result = rewriter.create<vector::InsertOp>(loc, el, result, i);
     }
   } else {
     if (auto vecType = dyn_cast<VectorType>(loadedElType)) {
@@ -832,7 +831,7 @@ createNonLdMatrixLoads(RewriterBase &rewriter, vector::TransferReadOp op,
         Value el = rewriter.create<memref::LoadOp>(op.getLoc(), loadedElType,
                                                    op.getSource(), newIndices);
         result = rewriter.create<vector::InsertOp>(
-            op.getLoc(), el, result, rewriter.getI64ArrayAttr({i, innerIdx}));
+            op.getLoc(), el, result, ArrayRef<int64_t>{i, innerIdx});
       }
     }
   }

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

@@ -1025,44 +1025,37 @@ class VectorExtractOpConversion
     auto loc = extractOp->getLoc();
     auto resultType = extractOp.getResult().getType();
     auto llvmResultType = typeConverter->convertType(resultType);
-    auto positionArrayAttr = extractOp.getPosition();
+    ArrayRef<int64_t> positionArray = extractOp.getPosition();
 
     // 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()) {
+    if (positionArray.empty()) {
       rewriter.replaceOp(extractOp, adaptor.getVector());
       return success();
     }
 
     // One-shot extraction of vector from array (only requires extractvalue).
     if (isa<VectorType>(resultType)) {
-      SmallVector<int64_t> indices;
-      for (auto idx : positionArrayAttr.getAsRange<IntegerAttr>())
-        indices.push_back(idx.getInt());
       Value extracted = rewriter.create<LLVM::ExtractValueOp>(
-          loc, adaptor.getVector(), indices);
+          loc, adaptor.getVector(), positionArray);
       rewriter.replaceOp(extractOp, extracted);
       return success();
     }
 
     // Potential extraction of 1-D vector from array.
     Value extracted = adaptor.getVector();
-    auto positionAttrs = positionArrayAttr.getValue();
-    if (positionAttrs.size() > 1) {
-      SmallVector<int64_t> nMinusOnePosition;
-      for (auto idx : positionAttrs.drop_back())
-        nMinusOnePosition.push_back(cast<IntegerAttr>(idx).getInt());
-      extracted = rewriter.create<LLVM::ExtractValueOp>(loc, extracted,
-                                                        nMinusOnePosition);
+    if (positionArray.size() > 1) {
+      extracted = rewriter.create<LLVM::ExtractValueOp>(
+          loc, extracted, positionArray.drop_back());
     }
 
     // Remaining extraction of element from 1-D LLVM vector
-    auto position = cast<IntegerAttr>(positionAttrs.back());
     auto i64Type = IntegerType::get(rewriter.getContext(), 64);
-    auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position);
+    auto constant =
+        rewriter.create<LLVM::ConstantOp>(loc, i64Type, positionArray.back());
     extracted =
         rewriter.create<LLVM::ExtractElementOp>(loc, extracted, constant);
     rewriter.replaceOp(extractOp, extracted);
@@ -1147,52 +1140,48 @@ class VectorInsertOpConversion
     auto sourceType = insertOp.getSourceType();
     auto destVectorType = insertOp.getDestVectorType();
     auto llvmResultType = typeConverter->convertType(destVectorType);
-    auto positionArrayAttr = insertOp.getPosition();
+    ArrayRef<int64_t> positionArray = insertOp.getPosition();
 
     // 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()) {
+    if (positionArray.empty()) {
       rewriter.replaceOp(insertOp, adaptor.getSource());
       return success();
     }
 
     // One-shot insertion of a vector into an array (only requires insertvalue).
     if (isa<VectorType>(sourceType)) {
       Value inserted = rewriter.create<LLVM::InsertValueOp>(
-          loc, adaptor.getDest(), adaptor.getSource(),
-          LLVM::convertArrayToIndices(positionArrayAttr));
+          loc, adaptor.getDest(), adaptor.getSource(), positionArray);
       rewriter.replaceOp(insertOp, inserted);
       return success();
     }
 
     // Potential extraction of 1-D vector from array.
     Value extracted = adaptor.getDest();
-    auto positionAttrs = positionArrayAttr.getValue();
-    auto position = cast<IntegerAttr>(positionAttrs.back());
     auto oneDVectorType = destVectorType;
-    if (positionAttrs.size() > 1) {
+    if (positionArray.size() > 1) {
       oneDVectorType = reducedVectorTypeBack(destVectorType);
       extracted = rewriter.create<LLVM::ExtractValueOp>(
-          loc, extracted,
-          LLVM::convertArrayToIndices(positionAttrs.drop_back()));
+          loc, extracted, positionArray.drop_back());
     }
 
     // Insertion of an element into a 1-D LLVM vector.
     auto i64Type = IntegerType::get(rewriter.getContext(), 64);
-    auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position);
+    auto constant =
+        rewriter.create<LLVM::ConstantOp>(loc, i64Type, positionArray.back());
     Value inserted = rewriter.create<LLVM::InsertElementOp>(
         loc, typeConverter->convertType(oneDVectorType), extracted,
         adaptor.getSource(), constant);
 
     // Potential insertion of resulting 1-D vector into array.
-    if (positionAttrs.size() > 1) {
+    if (positionArray.size() > 1) {
       inserted = rewriter.create<LLVM::InsertValueOp>(
-          loc, adaptor.getDest(), inserted,
-          LLVM::convertArrayToIndices(positionAttrs.drop_back()));
+          loc, adaptor.getDest(), inserted, positionArray.drop_back());
     }
 
     rewriter.replaceOp(insertOp, inserted);

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

@@ -886,10 +886,9 @@ struct UnrollTransferReadConversion
   /// vector::InsertOp, return that operation's indices.
   void getInsertionIndices(TransferReadOp xferOp,
                            SmallVector<int64_t, 8> &indices) const {
-    if (auto insertOp = getInsertOp(xferOp)) {
-      for (Attribute attr : insertOp.getPosition())
-        indices.push_back(dyn_cast<IntegerAttr>(attr).getInt());
-    }
+    if (auto insertOp = getInsertOp(xferOp))
+      indices.assign(insertOp.getPosition().begin(),
+                     insertOp.getPosition().end());
   }
 
   /// Rewrite the op: Unpack one dimension. Can handle masks, out-of-bounds
@@ -1013,10 +1012,9 @@ struct UnrollTransferWriteConversion
   /// indices.
   void getExtractionIndices(TransferWriteOp xferOp,
                             SmallVector<int64_t, 8> &indices) const {
-    if (auto extractOp = getExtractOp(xferOp)) {
-      for (Attribute attr : extractOp.getPosition())
-        indices.push_back(dyn_cast<IntegerAttr>(attr).getInt());
-    }
+    if (auto extractOp = getExtractOp(xferOp))
+      indices.assign(extractOp.getPosition().begin(),
+                     extractOp.getPosition().end());
   }
 
   /// Rewrite the op: Unpack one dimension. Can handle masks, out-of-bounds

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

@@ -152,7 +152,7 @@ struct VectorExtractOpConvert final
       return success();
     }
 
-    int32_t id = getFirstIntValue(extractOp.getPosition());
+    int32_t id = extractOp.getPosition()[0];
     rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
         extractOp, adaptor.getVector(), id);
     return success();
@@ -232,7 +232,7 @@ struct VectorInsertOpConvert final
       return success();
     }
 
-    int32_t id = getFirstIntValue(insertOp.getPosition());
+    int32_t id = insertOp.getPosition()[0];
     rewriter.replaceOpWithNewOp<spirv::CompositeInsertOp>(
         insertOp, adaptor.getSource(), adaptor.getDest(), id);
     return success();