[mlir][vector] Add constant folder for vector.shuffle ops

[ThomasRaoux]

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

1830

nit: type shouldn't change. Can you just use getVectorType() instead?

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

%1 = vector.shuffle %c, %b[0, 1, 2]
                : vector<2x16xf32>, vector<1x16xf32> ; yields vector<3x16xf32>
     %2 = vector.shuffle %a, %b[3, 2, 1, 0]
                : vector<2xf32>, vector<2xf32>       ; yields vector<4xf32>
     ```
   }];
   let builders = [
     OpBuilder<(ins "Value":$v1, "Value":$v2, "ArrayRef<int64_t>")>
   ];
+  let hasFolder = 1;
   let extraClassDeclaration = [{
     static StringRef getMaskAttrName() { return "mask"; }
     VectorType getV1VectorType() {
       return v1().getType().cast<VectorType>();
     }
     VectorType getV2VectorType() {
       return v2().getType().cast<VectorType>();
     }

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

static ParseResult parseShuffleOp(OpAsmParser &parser, OperationState &result) {
   shape.push_back(maskLength);
   for (int64_t r = 1; r < v1Rank; ++r)
     shape.push_back(v1Type.getDimSize(r));
   VectorType resType = VectorType::get(shape, v1Type.getElementType());
   parser.addTypeToList(resType, result.types);
   return success();
 }
 
+OpFoldResult vector::ShuffleOp::fold(ArrayRef<Attribute> operands) {
+  Attribute lhs = operands.front(), rhs = operands.back();
+  if (!lhs || !rhs)
+    return {};
+
+  auto lhsType = lhs.getType().cast<VectorType>();
+  // Only support 1-D for now to avoid complicated n-D DenseElementsAttr
+  // manipulation.
+  if (lhsType.getRank() != 1)
+    return {};
+  int64_t lhsSize = lhsType.getDimSize(0);
+
+  SmallVector<Attribute> results;
+  auto lhsElements = lhs.cast<DenseElementsAttr>().getValues<Attribute>();
+  auto rhsElements = rhs.cast<DenseElementsAttr>().getValues<Attribute>();
+  for (const auto &index : this->mask().getAsValueRange<IntegerAttr>()) {
+    int64_t i = index.getZExtValue();
+    if (i >= lhsSize) {
+      results.push_back(rhsElements[i - lhsSize]);
+    } else {
+      results.push_back(lhsElements[i]);
+    }
+  }
+
+  return DenseElementsAttr::get(getVectorType(), results);
+}
+
 //===----------------------------------------------------------------------===//
 // InsertElementOp
 //===----------------------------------------------------------------------===//
 
 void InsertElementOp::build(OpBuilder &builder, OperationState &result,
                             Value source, Value dest) {
   result.addOperands({source, dest});
   result.addTypes(dest.getType());

mlir/test/Dialect/Vector/canonicalize.mlir

 func @splat_fold() -> vector<4xf32> {
   %c = arith.constant 1.0 : f32
   %v = vector.splat %c : vector<4xf32>
   return %v : vector<4xf32>
 
   // CHECK-NEXT: [[V:%.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
   // CHECK-NEXT: return [[V]] : vector<4xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @shuffle_1d
+//       CHECK:   %[[V:.+]] = arith.constant dense<[3, 2, 5, 1]> : vector<4xi32>
+//       CHECK:   return %[[V]]
+func @shuffle_1d() -> vector<4xi32> {
+  %v0 = arith.constant dense<[0, 1, 2]> : vector<3xi32>
+  %v1 = arith.constant dense<[3, 4, 5]> : vector<3xi32>
+  %shuffle = vector.shuffle %v0, %v1 [3, 2, 5, 1] : vector<3xi32>, vector<3xi32>
+  return %shuffle : vector<4xi32>
+}