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

[Diego Caballero]

================
@@ -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();
+  }
----------------
dcaballe wrote:

Don't we have a interface for these two ops that we can use here?

If the implementation (below) is different we may want to consider having two methods, one for pack, one for unpack.

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