[mlir] [VectorOps] Implement vector tuple get folding

[aartbik]

Summary

Rewrites get-i tup<a1,...,an> into ai

[merge_guards_bot]

Unit tests: fail. 62015 tests passed, 2 failed and 783 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_mutex_requirements_mutex/thread_mutex_recursive/lock.pass.cpp
failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_mutex_requirements_mutex/thread_mutex_recursive/try_lock.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 62017 tests passed, 0 failed and 783 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[andydavis1]

Looks good. Thanks Aart!

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1696

I see. So the idea is to rewrite users one at a time, and eventually the tuple op will have no more users and get canonicalized away?

mlir/test/Dialect/VectorOps/vector-transforms.mlir

306

Is it worth adding a test which folds away all of the tuple op users, and test that the tuple op gets canonicalized away?

[aartbik]

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1696

Correct! A subsequent CL will refactor the vector-slices transformations as a standalone pass, so we can test the intermediate steps better!

mlir/test/Dialect/VectorOps/vector-transforms.mlir

306

I have a couple more CLs coming up with this.

[rriddle]

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1686

Like I said before this should be a fold.

[aartbik]

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1686

I had added a question in the prior CL. The tuple is not a constant, and I was following the pattern that was there wrt using a canonicalization.

[aartbik]

folder

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1686

One folder coming up for Monsieur!

[merge_guards_bot]

Unit tests: pass. 62017 tests passed, 0 failed and 783 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[aartbik]

Is this good to go? I have two follow up CLs in the pipeline!

[nicolasvasilache]

mlir/lib/Dialect/VectorOps/VectorOps.cpp

1691

drop trivial braces plz

[aartbik]

drop the braces, it is cleaner

[merge_guards_bot]

Unit tests: fail. 62132 tests passed, 5 failed and 807 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

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

let extraClassDeclaration = [{
     VectorType getResultVectorType() {
       return getResult().getType().cast<VectorType>();
     }
     int64_t getIndex() {
       return getAttrOfType<IntegerAttr>("index").getValue().getSExtValue();
     }
     static StringRef getIndexAttrName() { return "index"; }
   }];
+  let hasFolder = 1;
 }
 
 def Vector_PrintOp :
   Vector_Op<"print", []>, Arguments<(ins AnyType:$source)> {
   let summary = "print operation (for testing and debugging)";
   let description = [{
     Prints the source vector (or scalar) to stdout in human readable
     format (for testing and debugging). No return value.

mlir/lib/Dialect/VectorOps/VectorOps.cpp

 static LogicalResult verify(TupleGetOp op) {
   auto tupleType = op.getOperand().getType().cast<TupleType>();
   if (op.getIndex() < 0 ||
       op.getIndex() >= static_cast<int64_t>(tupleType.size()))
     return op.emitOpError("tuple get index out of range");
   return success();
 }
 
+OpFoldResult TupleGetOp::fold(ArrayRef<Attribute> operands) {
+  // Rewrite:
+  //    %t = vector.tuple .., %e_i, ..
+  //    %x = vector.tuple_get %t, i
+  // into:
+  //    %t = vector.tuple .., %e_i, ..  // one less use
+  //    %x = %e_i
+  if (auto tupleOp = dyn_cast_or_null<TupleOp>(getOperand().getDefiningOp()))
+    return tupleOp.getOperand(getIndex());
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantMaskOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseConstantMaskOp(OpAsmParser &parser,
                                        OperationState &result) {
   Type resultType;
   ArrayAttr maskDimSizesAttr;
   StringRef attrName = ConstantMaskOp::getMaskDimSizesAttrName();

mlir/test/Dialect/VectorOps/vector-transforms.mlir

affine.for %arg3 = 0 to %arg1 step 4 {
       %4 = vector.transfer_read %0[%arg2, %arg3], %cst  {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : memref<?x?xf32>, vector<4x4xf32>
       %5 = vector.transfer_read %1[%arg2, %arg3], %cst  {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : memref<?x?xf32>, vector<4x4xf32>
       %6 = addf %4, %5 : vector<4x4xf32>
       vector.transfer_write %6, %2[%arg2, %arg3] {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : vector<4x4xf32>, memref<?x?xf32>
     }
   }
   return
 }
+
+// CHECK-LABEL: func @tuple_get(%arg0: vector<4xf32>, %arg1: vector<8xf32>)
+//       CHECK: return %arg1
+
+func @tuple_get(%arg0: vector<4xf32>, %arg1: vector<8xf32>) -> vector<8xf32> {
+  %0 = vector.tuple %arg0, %arg1 : vector<4xf32>, vector<8xf32>
+  %1 = vector.tuple_get %0, 1 : tuple<vector<4xf32>, vector<8xf32>>
+  return %1 : vector<8xf32>
+}