[mlir][linalg] Vectorize unpack op without masking #89067

pashu123 · 2024-04-17T12:56:23Z

Enables vectorization of unpack op in the case of unknown vector size.
The vector sizes are determined by the result's shape.

llvmbot · 2024-04-17T12:56:55Z

@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Prashant Kumar (pashu123)

Changes

…t vector size

In case, the vector sizes are not provided for the vectorization of tensor.unpack op, the vector sizes are determined by the result shape. This also assumes that the input and output shapes are static.

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

2 Files Affected:

(modified) mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp (+20-3)
(modified) mlir/test/Dialect/Linalg/vectorization.mlir (+23)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index df61381432921b..92d2d129ff749c 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -1597,6 +1597,16 @@ vectorizeAsTensorUnpackOp(RewriterBase &rewriter, tensor::UnPackOp unpackOp,
 
   RankedTensorType unpackTensorType = unpackOp.getSourceType();
 
+  // If the input vector sizes are not provided, then the vector sizes are
+  // determined by the result tensor shape. In case the vector sizes aren't
+  // provided, we update the inBounds attribute instead of masking.
+  bool doMasking = true;
+  if (inputVectorSizes.empty()) {
+    ArrayRef<int64_t> resultTensorShape = unpackOp.getDestType().getShape();
+    inputVectorSizes = resultTensorShape.take_front(unpackOp.getSourceRank());
+    doMasking = false;
+  }
+
   ArrayRef<int64_t> innerDimPos = unpackOp.getInnerDimsPos();
   ArrayRef<int64_t> innerTiles = unpackOp.getStaticInnerTiles();
 
@@ -1651,7 +1661,8 @@ vectorizeAsTensorUnpackOp(RewriterBase &rewriter, tensor::UnPackOp unpackOp,
   // to shape of source, then a mask is necessary.
   Value readResult = createReadOrMaskedRead(
       rewriter, loc, unpackOp.getSource(),
-      ArrayRef<int64_t>(readMaskShape.begin(), readMaskShape.end()), padValue);
+      ArrayRef<int64_t>(readMaskShape.begin(), readMaskShape.end()), padValue,
+      doMasking);
 
   PackingMetadata packMetadata;
   SmallVector<int64_t> lastDimToInsertPosPerm =
@@ -1827,8 +1838,14 @@ vectorizeUnPackOpPrecondition(tensor::UnPackOp unpackOp,
     LDBG("Inner-tiles must be constant: " << unpackOp << "\n");
     return failure();
   }
-  llvm::ArrayRef<int64_t> resultShape = unpackOp.getDestType().getShape();
-  if (!inputVectorSizes.empty() &&
+  ArrayRef<int64_t> resultShape = unpackOp.getDestType().getShape();
+  bool satisfyEmptyCond = true;
+  if (inputVectorSizes.empty()) {
+    if (!unpackOp.getDestType().hasStaticShape() ||
+        !unpackOp.getSourceType().hasStaticShape())
+      satisfyEmptyCond = false;
+  }
+  if (!satisfyEmptyCond &&
       failed(isValidMaskedInputVector(resultShape, inputVectorSizes)))
     return failure();
 
diff --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir
index 80a5a4c6702ac1..5a81853973906b 100644
--- a/mlir/test/Dialect/Linalg/vectorization.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization.mlir
@@ -985,3 +985,26 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+  // -----
+
+func.func @test_vectorize_unpack_no_vector_sizes(%source: tensor<8x8x32x16xf32>, %dest: tensor<256x128xf32>) -> tensor<256x128xf32> {
+  // CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: %[[READ:.*]] = vector.transfer_read {{.*}} : tensor<8x8x32x16xf32>, vector<8x8x32x16xf32>
+  // CHECK: %[[TRANSP:.*]] = vector.transpose %[[READ]], [0, 2, 1, 3] : vector<8x8x32x16xf32> to vector<8x32x8x16xf32>
+  // CHECK: %[[SHAPC:.*]] = vector.shape_cast %[[TRANSP]] : vector<8x32x8x16xf32> to vector<256x128xf32>
+  // CHECK: %[[EMPT:.*]] = tensor.empty() : tensor<256x128xf32>
+  // CHECK: %[[C00:.*]] = arith.constant 0 : index
+  // CHECK: %[[WRIT:.*]] = vector.transfer_write %[[SHAPC]], {{.*}} : vector<256x128xf32>, tensor<256x128xf32>
+  // CHECK: return %[[WRIT]] : tensor<256x128xf32>
+   %0 = tensor.unpack %source inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %dest : tensor<8x8x32x16xf32> -> tensor<256x128xf32>
+   return %0 : tensor<256x128xf32>
+ }
+ module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["tensor.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+   transform.structured.vectorize %0 : !transform.any_op
+    transform.yield
+  } 
+ }

pashu123 · 2024-04-17T12:57:36Z

@hanhanW I am in the process of adding more tests. Thanks.

banach-space · 2024-04-17T14:05:15Z

[nit] Could you trim the commit subject? It's much easier to read if it fits in one line. Also:

https://mlir.llvm.org/getting_started/Contributing/#commit-messages --> https://cbea.ms/git-commit/#limit-50

Thanks :)

pashu123 · 2024-04-17T14:59:36Z

[nit] Could you trim the commit subject? It's much easier to read if it fits in one line. Also:

https://mlir.llvm.org/getting_started/Contributing/#commit-messages --> https://cbea.ms/git-commit/#limit-50

Thanks :)

Thanks for the reference. I have updated the message and the body. PTAL.

hanhanW · 2024-04-18T20:44:38Z

mlir/test/Dialect/Linalg/vectorization.mlir

We need a test for unpack which also slices output. E.g.,

%0 = tensor.unpack %source inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %dest : tensor<8x8x32x16xf32> -> tensor<255x127xf32>

Would the vector sizes for this case be inner_tiles[x] * source_dim[inner_dims_pos[x]] for x in len(inner_tiles) and then the inbounds will be set accordingly?

Yes, I think so. The inbounds (of the xfer_write op) will be set accordingly.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

hanhanW · 2024-04-25T16:53:22Z

mlir/test/Dialect/Linalg/vectorization.mlir

Yes, I think so. The inbounds (of the xfer_write op) will be set accordingly.

hanhanW

thanks!

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

banach-space

Please could you address my comments before landing this? It feels that there's scope for better code re-use.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

banach-space · 2024-05-01T10:08:35Z

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

+  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;
+  }


Is yet another variable, initVectorShape, really needed? Also, see my comment below.

Suggested change

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;

}

if (inputVectorSizes.empty()) {

inputVectorSizes.resize(sourceShape.take_front(destSize))

if (!outerDimsPerm.empty())

applyPermutationToVector(inputVectorSizes, outerDimsPerm);

for (auto [i, pos] : llvm::enumerate(innerDimPos))

inputVectorSizes[pos] *= innerTiles[i];

useInBoundsInsteadOfMasking = true;

}

We won't be able to modify/resize the inputVectorShape var since it's an ArrayRef, but we can point it to another variable. In that sense, it's needed.

I think we can't do the trick because it is ArrayRef type. We can only apply permutation on SmallVector.

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:

// 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).

I've tried to simplify it so we don't have the else block. Also, we can't use vectorSizes = inputVectorSizes SmallVectors can't point to ArrayRefs but the other way around is possible.

banach-space · 2024-05-01T10:15:36Z

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

  SmallVector<int64_t> readMaskShape(inputVectorSizes.begin(),
                                     inputVectorSizes.end());


At this point useInBoundsInsteadOfMasking is already set - why do we bother defining and calculating readMaskShape if useInBoundsInsteadOfMasking is true? It feels like something that should be factored out to a dedicated hook, e.g. computeReadMaskShapeForUnpackOp, and then:

SmallVector<int64_t> readMaskShape; if (!useInBoundsInsteadOfMasking) readMaskShape = computeReadMaskShapeForUnpackOp();

Also, it looks like the shape calculation for readMaskShape and initVectorShape are duplicated? Again, why not introduce a dedicated hook for that?

Well, I think readMaskShape is not named properly. It's actually readVectorSizes and we have to unconditionally define it.

Good point, I missed that!

readMaskShape is used when calling createReadOrMaskedRead:

llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Lines 1613 to 1616 in 6d44a1e

Value readResult = vector::createReadOrMaskedRead(

rewriter, loc, unpackOp.getSource(),

ArrayRef<int64_t>(readMaskShape.begin(), readMaskShape.end()), padValue,

/*useInBoundsInsteadOfMasking=*/false);

And, the signature of createReadOrMaskedRead is here:

llvm-project/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

Lines 195 to 197 in 6d44a1e

Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,

ArrayRef<int64_t> readShape, Value padValue,

bool useInBoundsInsteadOfMasking);

So it's not really readMaskShape, it's readShape or readVectorSizes like you suggested. Am I correct that we only need to calculate this once?

Yes, we need to calculate it once. I've renamed this to readVectorSizes since the mask is generated by createReadOrMaskedRead based on the vectorSizes. The same goes for writeMaskShape.

mlir/test/Dialect/Linalg/vectorization.mlir

github-actions · 2024-05-02T12:19:05Z

✅ With the latest revision this PR passed the C/C++ code formatter.

banach-space

LGTM, modulo a few minor suggestions, thanks!

This is looking much better now and is much easier to follow, thank you @pashu123 !

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

banach-space · 2024-05-03T14:36:28Z

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

+  // vectorSizes is the shape of the vector that will be used to do final
+  // write on the destination 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:
+  // - vectorSizes = 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.


Suggested change

// vectorSizes is the shape of the vector that will be used to do final

// write on the destination 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:

// - vectorSizes = 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.

// vectorSizes is the shape of the vector that will be used to do final

// write on the destination tensor. It is set like this: Let's say the

// source tensor is rank 'M' and the dest tensor is rank 'N', where N <= M.

// Thus:

// 1. vectorSizes = sourceShape.take_front(N)

// 2. if outer_dims_perms is present: do that permutation on vectorSizes.

// 3. multiply all the locations in VectorSize pointed by innerDimPos by the innerTiles

// attribute value.

Remove references to initVectorShape

This sentence doesn't make sense when vectorSizes is empty: "the vectorSize (VS) array is size 'N'". That's fine - I think what's more important is the rank of the source tensor (M) and the output tensor (N).

Consistent capitalisation (nit)

Use numbering to highlight that these are consecutive steps (nit)

Done. Thanks. I need to improve my comments.

banach-space · 2024-05-03T14:40:04Z

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

+  // - 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> vectorSizes(inputVectorSizes);


Sounds like vectorSizes could be renamed as writeVectorSizes?

If !inputVectorSizes.empty(), add assert(inputVectorSizes.size() == destSize && "Incorrect number of input vector sizes"); (unless I got this one wrong?)

There's actually a check performed here:

llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Line 1676 in 64013f9

SmallVector<int64_t> writeVectorSizes(

. Only if the destination type is static can we use vectorSizes; otherwise, we resort to something else.

check is performed here:

llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Line 1798 in 64013f9

if (!satisfyEmptyCond &&

Yeah, looks like this condition is indeed checked in 2. above, thanks!

That's a "pre-condition" though - no harm in adding an additional assert to document assumptions made in this method.

In any case, it's just a nice to have :)

Added. thanks.

Thanks again for working on this - that's greatly appreciated 🙏🏻

Enables vectorization of unpack op in the case of unknown vector size. The vector sizes are determined by the result shape.

pashu123 requested review from hanhanW and nicolasvasilache as code owners April 17, 2024 12:56

llvmbot added mlir:linalg mlir labels Apr 17, 2024

pashu123 force-pushed the unpack_nosize branch from 356a4ab to 8e7d415 Compare April 17, 2024 14:57

pashu123 changed the title ~~Add support for static unpack op vectorization without providing inpu…~~ [mlir] Vectorize unpack op given no vector sizes Apr 17, 2024

hanhanW changed the title ~~[mlir] Vectorize unpack op given no vector sizes~~ [mlir][linalg] Vectorize unpack op without masking Apr 17, 2024

hanhanW requested changes Apr 18, 2024

View reviewed changes

hanhanW requested a review from chelini April 18, 2024 20:44

banach-space reviewed Apr 19, 2024

View reviewed changes

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp Outdated Show resolved Hide resolved

pashu123 force-pushed the unpack_nosize branch 2 times, most recently from 5bc4819 to 1915d7e Compare April 25, 2024 11:53

hanhanW reviewed Apr 25, 2024

View reviewed changes

pashu123 force-pushed the unpack_nosize branch 3 times, most recently from ed8e2a6 to eb5f6ee Compare April 30, 2024 14:47

pashu123 requested a review from hanhanW April 30, 2024 18:17

hanhanW approved these changes Apr 30, 2024

View reviewed changes

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp Outdated Show resolved Hide resolved

pashu123 force-pushed the unpack_nosize branch from eb5f6ee to 6332c1b Compare May 1, 2024 09:33

banach-space requested changes May 1, 2024

View reviewed changes

pashu123 force-pushed the unpack_nosize branch from 6332c1b to f491235 Compare May 2, 2024 12:15

pashu123 force-pushed the unpack_nosize branch 3 times, most recently from 17ec14f to 64013f9 Compare May 3, 2024 10:57

banach-space approved these changes May 3, 2024

View reviewed changes

pashu123 force-pushed the unpack_nosize branch from 64013f9 to 1eda5b8 Compare May 3, 2024 15:07

[mlir] Vectorize unpack op given no vector sizes

7e14b29

Enables vectorization of unpack op in the case of unknown vector size. The vector sizes are determined by the result shape.

pashu123 force-pushed the unpack_nosize branch from 1eda5b8 to 7e14b29 Compare May 3, 2024 15:38

pashu123 merged commit 2755c69 into llvm:main May 3, 2024

		SmallVector<int64_t> readMaskShape(inputVectorSizes.begin(),
		inputVectorSizes.end());

	Value readResult = vector::createReadOrMaskedRead(
	rewriter, loc, unpackOp.getSource(),
	ArrayRef<int64_t>(readMaskShape.begin(), readMaskShape.end()), padValue,
	/useInBoundsInsteadOfMasking=/false);

	Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,
	ArrayRef<int64_t> readShape, Value padValue,
	bool useInBoundsInsteadOfMasking);

[mlir][linalg] Vectorize unpack op without masking #89067

[mlir][linalg] Vectorize unpack op without masking #89067

Uh oh!

Conversation

pashu123 commented Apr 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Apr 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pashu123 commented Apr 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

banach-space commented Apr 17, 2024

Uh oh!

pashu123 commented Apr 17, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hanhanW left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented May 2, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

pashu123 commented Apr 17, 2024 •

edited

Loading

llvmbot commented Apr 17, 2024 •

edited

Loading

pashu123 commented Apr 17, 2024 •

edited

Loading

github-actions bot commented May 2, 2024 •

edited

Loading