[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:

Still feels like this should be cleaner with two versions, one for packOp and one for unpackOp. The shared code seems to be only 106-111. Perhaps we could move that to a helper function and call it from the two versions?

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