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

[Aart Bik via Phabricator via llvm-commits]

aartbik created this revision.
Herald added subscribers: llvm-commits, liufengdb, lucyrfox, mgester, arpith-jacob, antiagainst, shauheen, burmako, jpienaar, rriddle, mehdi_amini.
Herald added a reviewer: nicolasvasilache.
Herald added a reviewer: nicolasvasilache.
Herald added a project: LLVM.
aartbik added reviewers: rriddle, andydavis1.

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


Repository:
  rG LLVM Github Monorepo

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,13 @@
   }
   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,36 @@
   return success();
 }
 
+namespace {
+
+class TupleGetFolder : public OpRewritePattern<TupleGetOp> {
+public:
+  using OpRewritePattern<TupleGetOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(TupleGetOp op,
+                                     PatternRewriter &rewriter) const override {
+    // 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>(op.getOperand().getDefiningOp())) {
+      rewriter.replaceOp(op, tupleOp.getOperand(op.getIndex()));
+      return matchSuccess();
+    }
+    return matchFailure();
+  }
+};
+
+} // end anonymous namespace
+
+void TupleGetOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                             MLIRContext *context) {
+  results.insert<TupleGetFolder>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantMaskOp
 //===----------------------------------------------------------------------===//
@@ -1814,7 +1844,9 @@
 
 void mlir::vector::populateVectorToVectorCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
-  patterns.insert<CreateMaskFolder, StridedSliceConstantMaskFolder>(context);
+  patterns
+      .insert<TupleGetFolder, CreateMaskFolder, StridedSliceConstantMaskFolder>(
+          context);
 }
 
 namespace mlir {
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 hasCanonicalizer = 1;
 }
 
 def Vector_PrintOp :


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