[Mlir-commits] [mlir] [mlir][linalg] Vectorize unpack op without masking (PR #89067)

[Andrzej Warzyński]

================
@@ -1560,11 +1574,32 @@ vectorizeAsTensorUnpackOp(RewriterBase &rewriter, tensor::UnPackOp unpackOp,
 
   ArrayRef<int64_t> innerDimPos = unpackOp.getInnerDimsPos();
   ArrayRef<int64_t> innerTiles = unpackOp.getStaticInnerTiles();
+  ArrayRef<int64_t> sourceShape = unpackTensorType.getShape();
+  bool useInBoundsInsteadOfMasking = false;
+  ArrayRef<int64_t> outerDimsPerm = unpackOp.getOuterDimsPerm();
+
+  auto destSize = unpackOp.getDestRank();
+
+  // initVectorShape is the shape of the vector that will be read from the
+  // source tensor. It is set like this: Let's say the sourceShape is 'M' and
+  // the vectorSize (VS) array is size 'N' where N <= M. Thus:
+  // - initVectorShape = sourceShape.take_front(N)
+  // - if outer_dims_perms is present: do that permutation on initVectorShape.
+  // - Multiply all the locations pointed by innerDimPos by the innerTileSize
+  //  attribute value.
+  SmallVector<int64_t> initVectorShape(sourceShape.take_front(destSize));
+  if (inputVectorSizes.empty()) {
+    if (!outerDimsPerm.empty())
+      applyPermutationToVector(initVectorShape, outerDimsPerm);
+    for (auto [i, pos] : llvm::enumerate(innerDimPos))
+      initVectorShape[pos] *= innerTiles[i];
+
+    inputVectorSizes = initVectorShape;
+    useInBoundsInsteadOfMasking = true;
+  }
----------------
banach-space wrote:

My bad, sorry and thanks for checking!

I would still consider refining a bit. Especially, given that with this change, **vector sizes might come from 2 different places**. Here's what I'd do:
```cpp
// Keep the signature as is
static LogicalResult
vectorizeAsTensorUnpackOp(RewriterBase &rewriter, tensor::UnPackOp unpackOp,
                          ArrayRef<int64_t> inputVectorSizes,
                          SmallVectorImpl<Value> &newResults) {
           
   SmallVector<int64_t> vectorSizes;
   if (!inputVectorSizes.empty()) {
      vectorSizes = inputVectorSizes;
   } else {
      vectorSizes.resize(sourceShape.take_front(destSize))
      // The logic that you have added
   }
   
   // Later in this method, use `vectorSizes` rather than `inputVectorSizes`
   Operation *write = createWriteOrMaskedWrite(
      rewriter, loc, maskedRead, reifiedReturnShapes[0], vectorSizes,
      /*useInBoundsInsteadOfMasking=*/false);
}
```

Basically, if `inputVectorSizes` is used everywhere then that's suggesting that it's always the input parameter (`inputVectorSizes `) that's used for defining the "vector sizes" to use. With this change, that's no longer the case.

This comment is a nit, feel free to ignore (naming is hard).

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