[PATCH] D73213: [mlir] [VectorOps] Implement vector tuple get folding

[Aart Bik via Phabricator via llvm-commits]

aartbik updated this revision to Diff 239669.
aartbik marked 2 inline comments as done.
aartbik added a comment.

folder


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D73213/new/

https://reviews.llvm.org/D73213

Files:
  mlir/include/mlir/Dialect/VectorOps/VectorOps.td
  mlir/lib/Dialect/VectorOps/VectorOps.cpp
  mlir/test/Dialect/VectorOps/vector-transforms.mlir


Index: mlir/test/Dialect/VectorOps/vector-transforms.mlir
===================================================================
--- mlir/test/Dialect/VectorOps/vector-transforms.mlir
+++ mlir/test/Dialect/VectorOps/vector-transforms.mlir
@@ -302,3 +302,12 @@
   }
   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>
+}
Index: mlir/lib/Dialect/VectorOps/VectorOps.cpp
===================================================================
--- mlir/lib/Dialect/VectorOps/VectorOps.cpp
+++ mlir/lib/Dialect/VectorOps/VectorOps.cpp
@@ -1681,6 +1681,19 @@
   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
 //===----------------------------------------------------------------------===//
Index: mlir/include/mlir/Dialect/VectorOps/VectorOps.td
===================================================================
--- mlir/include/mlir/Dialect/VectorOps/VectorOps.td
+++ mlir/include/mlir/Dialect/VectorOps/VectorOps.td
@@ -1115,6 +1115,7 @@
     }
     static StringRef getIndexAttrName() { return "index"; }
   }];
+  let hasFolder = 1;
 }
 
 def Vector_PrintOp :


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D73213.239669.patch
Type: text/x-patch
Size: 1838 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200122/21c07a8a/attachment.bin>