[mlir] [VectorOps] Implement vector tuple get folding

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

Reviewers: nicolasvasilache, rriddle, andydavis1

Reviewed By: nicolasvasilache, rriddle, andydavis1

Subscribers: merge_guards_bot, mehdi_amini, rriddle, jpienaar, burmako, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, liufengdb, llvm-commits

Tags: #llvm

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

Author: aartbik

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

@@ -1115,6 +1115,7 @@ def Vector_TupleGetOp :
     }
     static StringRef getIndexAttrName() { return "index"; }
   }];
+  let hasFolder = 1;
 }
 
 def Vector_PrintOp :

mlir/lib/Dialect/VectorOps/VectorOps.cpp

@@ -1681,6 +1681,18 @@ static LogicalResult verify(TupleGetOp op) {
   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
 //===----------------------------------------------------------------------===//

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

@@ -302,3 +302,12 @@ func @vector_transfers(%arg0: index, %arg1: index) {
   }
   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>
+}