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

[Balaji V. Iyer.]

================
@@ -73,33 +73,52 @@ mlir::tensor::computeTransposedType(RankedTensorType rankedTensorType,
   return transposedTensorType;
 }
 
-SmallVector<int64_t>
-mlir::tensor::getPackInverseDestPermutation(PackOp packOp) {
+SmallVector<int64_t> mlir::tensor::getPackUnPackInverseDestPerm(
+    std::variant<tensor::PackOp, tensor::UnPackOp> op) {
+  PackingMetadata pMetaData;
+  return getPackUnPackInverseDestPerm(op, pMetaData);
+}
+
+SmallVector<int64_t> mlir::tensor::getPackUnPackInverseDestPerm(
+    std::variant<tensor::PackOp, tensor::UnPackOp> op,
+    PackingMetadata &packingMetadata) {
+
+  llvm::ArrayRef<int64_t> innerDimsPos, outerPerm;
+  int64_t rank = 0;
+  bool isPackOp = std::holds_alternative<tensor::PackOp>(op);
+  if (isPackOp) {
+    tensor::PackOp packOp = std::get<tensor::PackOp>(op);
+    innerDimsPos = packOp.getInnerDimsPos();
+    rank = packOp.getDestType().getRank();
+    outerPerm = packOp.getOuterDimsPerm();
+  } else {
+    tensor::UnPackOp unpackOp = std::get<tensor::UnPackOp>(op);
+    innerDimsPos = unpackOp.getInnerDimsPos();
+    rank = unpackOp.getSourceType().getRank();
+    outerPerm = unpackOp.getOuterDimsPerm();
+  }
----------------
bviyer wrote:

Well the workings is dependent on whether it is a packOp or unpackOp. For example, for unPackOp we need to look at source rank whereas for the packOp we need the destination Rank. ALso, we need to look at outerPositions for packed But not for unpacked. However, 80% of the code is the same, this is why I decided to use the std::variant.

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