[Mlir-commits] [mlir] [mlir][Vectorizer] Added support to Vectorize tensor.unpack (PR #76087)

Balaji V. Iyer. llvmlistbot at llvm.org
Thu Feb 15 17:49:42 PST 2024


================
@@ -1559,6 +1571,90 @@ vectorizeAsTensorPackOp(RewriterBase &rewriter, tensor::PackOp packOp,
   return success();
 }
 
+/// Vectorize a `tensor::UnPackOp` without OuterDimsPerms to these 4 Ops:
+///   Vector::TransferReadOp - Reads a vector from the source tensor
+///   vector::TransposeOp - Transpose the Source tensor
+///   ShapeCastOp - Reshape the data based on the target.
+///   vector::TransferWriteOp. - Write the result vector back to the destination
+///   tensor
+static LogicalResult vectorizeAsUnpackOp(RewriterBase &rewriter,
+                                         tensor::UnPackOp unpackOp,
+                                         ArrayRef<int64_t> inputVectorSizes,
+                                         SmallVectorImpl<Value> &newResults) {
+
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(unpackOp);
+
+  RankedTensorType unpackTensorType = unpackOp.getSourceType();
+
+  SmallVector<int64_t> readMaskShape(unpackTensorType.getShape());
----------------
bviyer wrote:

this [c7ed75e](https://github.com/llvm/llvm-project/pull/76087/commits/c7ed75e39f79fdf8c4de880c7ea8d1800be347d4) should support outer_dims_perm, however I think your steps 2 and 3 should be switched because in the manual it says outer_dims_perm should be done before combining any other dimensions.

https://github.com/llvm/llvm-project/pull/76087


More information about the Mlir-commits mailing list