[mlir][Vector] Add constant folder for insertelement.

[jacquesguan]

Summary

This revision adds constant folder for vector.insertelement.

[ThomasRaoux]

please include file context on the reviews (meaning use -U9999999 when creating the patch): https://llvm.org/docs/Phabricator.html#id8)

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

1843

nit: breakup this line to match the style of the file/project.

[jacquesguan]

Rebase and address comment.

[jacquesguan]

In D122721#3419715, @ThomasRaoux wrote:

please include file context on the reviews (meaning use -U9999999 when creating the patch): https://llvm.org/docs/Phabricator.html#id8)

I am using the Arcanist to upload the revision, does this miss something? I will recheck the configuration. Thanks.

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

1843

Done, thanks.

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

def Vector_InsertElementOp :
   ];
   let extraClassDeclaration = [{
     Type getSourceType() { return getSource().getType(); }
     VectorType getDestVectorType() {
       return getDest().getType().cast<VectorType>();
     }
   }];
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 def Vector_InsertOp :
   Vector_Op<"insert", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      AllTypesMatch<["dest", "res"]>]>,
      Arguments<(ins AnyType:$source, AnyVector:$dest, I64ArrayAttr:$position)>,

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

LogicalResult InsertElementOp::verify() {
   }
   if (dstVectorType.getRank() != 1)
     return emitOpError("unexpected >1 vector rank");
   if (!getPosition())
     return emitOpError("expected position for 1-D vector");
   return success();
 }
 
+OpFoldResult vector::InsertElementOp::fold(ArrayRef<Attribute> operands) {
+  // Skip the 0-D vector here.
+  if (operands.size() < 3)
+    return {};
+
+  Attribute src = operands[0];
+  Attribute dst = operands[1];
+  Attribute pos = operands[2];
+  if (!src || !dst || !pos)
+    return {};
+
+  auto dstElements = dst.cast<DenseElementsAttr>().getValues<Attribute>();
+
+  SmallVector<Attribute> results(dstElements);
+
+  auto attr = pos.dyn_cast<IntegerAttr>();
+  uint64_t posIdx = attr.getInt();
+
+  results[posIdx] = src;
+
+  return DenseElementsAttr::get(getDestVectorType(), results);
+}
+
 //===----------------------------------------------------------------------===//
 // InsertOp
 //===----------------------------------------------------------------------===//
 
 void InsertOp::build(OpBuilder &builder, OperationState &result, Value source,
                      Value dest, ArrayRef<int64_t> position) {
   result.addOperands({source, dest});
   auto positionAttr = getVectorSubscriptAttr(builder, position);

mlir/test/Dialect/Vector/canonicalize.mlir

 
 // CHECK-LABEL: func @shuffle_nofold2
 //       CHECK:   %[[V:.+]] = vector.shuffle %arg0, %arg1 [0, 1, 2, 3] : vector<[4]xi32>, vector<[2]xi32>
 //       CHECK:   return %[[V]]
 func @shuffle_nofold2(%v0 : vector<[4]xi32>, %v1 : vector<[2]xi32>) -> vector<4xi32> {
   %shuffle = vector.shuffle %v0, %v1 [0, 1, 2, 3] : vector<[4]xi32>, vector<[2]xi32>
   return %shuffle : vector<4xi32>
 }
+
+// -----
+
+// CHECK-LABEL: func @insert_element_fold
+//       CHECK:   %[[V:.+]] = arith.constant dense<[0, 1, 7, 3]> : vector<4xi32>
+//       CHECK:   return %[[V]]
+func @insert_element_fold() -> vector<4xi32> {
+  %v = arith.constant dense<[0, 1, 2, 3]> : vector<4xi32>
+  %s = arith.constant 7 : i32
+  %i = arith.constant 2 : i32
+  %1 = vector.insertelement %s, %v[%i : i32] : vector<4xi32>
+  return %1 : vector<4xi32>
+}