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

[jacquesguan]

Summary

This revision adds constant folder for vector.extractelement.

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

let builders = [
     OpBuilder<(ins "Value":$source, "Value":$position)>,
   ];
   let extraClassDeclaration = [{
     VectorType getVectorType() {
       return getVector().getType().cast<VectorType>();
     }
   }];
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 def Vector_ExtractOp :
   Vector_Op<"extract", [NoSideEffect,
      PredOpTrait<"operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      DeclareOpInterfaceMethods<InferTypeOpInterface>]>,
     Arguments<(ins AnyVector:$vector, I64ArrayAttr:$position)>,

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

LogicalResult vector::ExtractElementOp::verify() {
   }
   if (vectorType.getRank() != 1)
     return emitOpError("unexpected >1 vector rank");
   if (!getPosition())
     return emitOpError("expected position for 1-D vector");
   return success();
 }
 
+OpFoldResult vector::ExtractElementOp::fold(ArrayRef<Attribute> operands) {
+  // Skip the 0-D vector here now.
+  if (operands.size() < 2)
+    return {};
+
+  Attribute src = operands[0];
+  Attribute pos = operands[1];
+  if (!src || !pos)
+    return {};
+
+  auto srcElements = src.cast<DenseElementsAttr>().getValues<Attribute>();
+
+  auto attr = pos.dyn_cast<IntegerAttr>();
+  uint64_t posIdx = attr.getInt();
+
+  return srcElements[posIdx];
+}
+
 //===----------------------------------------------------------------------===//
 // ExtractOp
 //===----------------------------------------------------------------------===//
 
 void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
                               Value source, ArrayRef<int64_t> position) {
   build(builder, result, source, getVectorSubscriptAttr(builder, position));
 }

mlir/test/Dialect/Vector/canonicalize.mlir

 //       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>
 }
+
+// -----
+
+// CHECK-LABEL: func @extract_element_fold
+//       CHECK:   %[[C:.+]] = arith.constant 5 : i32
+//       CHECK:   return %[[C]]
+func @extract_element_fold() -> i32 {
+  %v = arith.constant dense<[1, 3, 5, 7]> : vector<4xi32>
+  %i = arith.constant 2 : i32
+  %1 = vector.extractelement %v[%i : i32] : vector<4xi32>
+  return %1 : i32
+}