[mlir][vector] Add canonicalization for extract of splat

[ThomasRaoux]

Summary

Make canonicalization working on broadcast also work on splat op.

[antiagainst]

mlir/lib/Dialect/Vector/VectorOps.cpp

2264

Formatting is a bit off here?

[ThomasRaoux]

Fix formatting

mlir/lib/Dialect/Vector/VectorOps.cpp

while (insertOp || transposeOp) {
     // `vector.insert` / `vector.transpose`.
     // Compute insert/transpose for the next iteration.
     insertOp = insertionDest.getDefiningOp<vector::InsertOp>();
     transposeOp = insertionDest.getDefiningOp<vector::TransposeOp>();
   }
   return Value();
 }
 
-/// Fold extractOp with scalar result coming from BroadcastOp.
+/// Fold extractOp with scalar result coming from BroadcastOp or SplatOp.
 static Value foldExtractFromBroadcast(ExtractOp extractOp) {
-  auto broadcastOp = extractOp.vector().getDefiningOp<vector::BroadcastOp>();
-  if (!broadcastOp)
+  Operation *defOp = extractOp.vector().getDefiningOp();
+  if (!defOp || !isa<vector::BroadcastOp, SplatOp>(defOp))
     return Value();
-  if (extractOp.getType() == broadcastOp.getSourceType())
-    return broadcastOp.source();
+  Value source = defOp->getOperand(0);
+  if (extractOp.getType() == source.getType())
+    return source;
   auto getRank = [](Type type) {
     return type.isa<VectorType>() ? type.cast<VectorType>().getRank() : 0;
   };
-  unsigned broadcasrSrcRank = getRank(broadcastOp.getSourceType());
+  unsigned broadcasrSrcRank = getRank(source.getType());
   unsigned extractResultRank = getRank(extractOp.getType());
   if (extractResultRank < broadcasrSrcRank) {
     auto extractPos = extractVector<int64_t>(extractOp.position());
     unsigned rankDiff = broadcasrSrcRank - extractResultRank;
     extractPos.erase(
         extractPos.begin(),
         std::next(extractPos.begin(), extractPos.size() - rankDiff));
-    extractOp.setOperand(broadcastOp.source());
+    extractOp.setOperand(source);
     // OpBuilder is only used as a helper to build an I64ArrayAttr.
     OpBuilder b(extractOp.getContext());
     extractOp->setAttr(ExtractOp::getPositionAttrName(),
                        b.getI64ArrayAttr(extractPos));
     return extractOp.getResult();
   }
   // TODO: In case the rank of the broadcast source is greater than the rank of
   // the extract result this can be combined into a new broadcast op. This needs
if (!lowerDimMatch) {
           getI64SubArray(op.sizes(), /* dropFront=*/rankDiff),
           getI64SubArray(op.strides(), /* dropFront=*/rankDiff));
     }
     rewriter.replaceOpWithNewOp<BroadcastOp>(op, op.getType(), source);
     return success();
   }
 };
 
+/// Pattern to rewrite an ExtractStridedSliceOp(SplatOp) to SplatOp.
+class StridedSliceSplat final : public OpRewritePattern<ExtractStridedSliceOp> {
+public:
+  using OpRewritePattern<ExtractStridedSliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ExtractStridedSliceOp op,
+                                PatternRewriter &rewriter) const override {
+    auto splat = op.vector().getDefiningOp<SplatOp>();
+    if (!splat)
+      return failure();
+    rewriter.replaceOpWithNewOp<SplatOp>(op, op.getType(), splat.input());
+    return success();
+  }
+};
+
 } // end anonymous namespace
 
 void ExtractStridedSliceOp::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
   // Pattern to rewrite a ExtractStridedSliceOp(ConstantMaskOp) ->
   // ConstantMaskOp and ExtractStridedSliceOp(ConstantOp) -> ConstantOp.
   results.add<StridedSliceConstantMaskFolder, StridedSliceConstantFolder,
-              StridedSliceBroadcast>(context);
+              StridedSliceBroadcast, StridedSliceSplat>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // TransferReadOp
 //===----------------------------------------------------------------------===//
 
 template <typename EmitFun>
 static LogicalResult verifyPermutationMap(AffineMap permutationMap,

mlir/test/Dialect/Vector/canonicalize.mlir

 func @fold_extract_broadcast(%a : f32) -> f32 {
   %b = vector.broadcast %a : f32 to vector<1x2x4xf32>
   %r = vector.extract %b[0, 1, 2] : vector<1x2x4xf32>
   return %r : f32
 }
 
 // -----
 
+// CHECK-LABEL: fold_extract_splat
+//  CHECK-SAME:   %[[A:.*]]: f32
+//       CHECK:   return %[[A]] : f32
+func @fold_extract_splat(%a : f32) -> f32 {
+  %b = splat %a : vector<1x2x4xf32>
+  %r = vector.extract %b[0, 1, 2] : vector<1x2x4xf32>
+  return %r : f32
+}
+
+// -----
+
 // CHECK-LABEL: fold_extract_broadcast_vector
 //  CHECK-SAME:   %[[A:.*]]: vector<4xf32>
 //       CHECK:   return %[[A]] : vector<4xf32>
 func @fold_extract_broadcast_vector(%a : vector<4xf32>) -> vector<4xf32> {
   %b = vector.broadcast %a : vector<4xf32> to vector<1x2x4xf32>
   %r = vector.extract %b[0, 1] : vector<1x2x4xf32>
   return %r : vector<4xf32>
 }

 
 // CHECK-LABEL: func @insert_strided_slice_full_range
 //  CHECK-SAME: %[[SOURCE:.+]]: vector<16x16xf16>, %{{.+}}: vector<16x16xf16>
 func @insert_strided_slice_full_range(%source: vector<16x16xf16>, %dest: vector<16x16xf16>) -> vector<16x16xf16> {
   %0 = vector.insert_strided_slice %source, %dest {offsets = [0, 0], strides = [1, 1]} : vector<16x16xf16> into vector<16x16xf16>
   // CHECK: return %[[SOURCE]]
   return %0: vector<16x16xf16>
 }
+
+// -----
+
+// CHECK-LABEL: extract_strided_splat
+//       CHECK:   %[[B:.*]] = splat %{{.*}} : vector<2x4xf16>
+//  CHECK-NEXT:   return %[[B]] : vector<2x4xf16>
+func @extract_strided_splat(%arg0: f16) -> vector<2x4xf16> {
+ %0 = splat %arg0 : vector<16x4xf16>
+ %1 = vector.extract_strided_slice %0
+  {offsets = [1, 0], sizes = [2, 4], strides = [1, 1]} :
+  vector<16x4xf16> to vector<2x4xf16>
+  return %1 : vector<2x4xf16>
+}