Skip to content

[mlir][vector] Refactor createWriteOrMaskedWrite #138137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
banach-space merged 2 commits into from
May 14, 2025
Merged

[mlir][vector] Refactor createWriteOrMaskedWrite #138137

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
18cbf37
[mlir][vector] Refactor `createWriteOrMaskedWrite`
banach-space May 1, 2025
2cd73e6
fixup! [mlir][vector] Refactor `createWriteOrMaskedWrite`
banach-space May 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 55 additions & 43 deletions mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1506,84 +1506,86 @@ static SmallVector<int64_t> getTiledPackShape(linalg::PackOp packOp,
return applyPermutation(destShape, linalg::getPackInverseDestPerm(packOp));
}

/// Creates a TransferWriteOp to write `input` into a newly initialized
/// output tensor.
/// Creates an optionally masked TransferWriteOp
///
/// Given:
/// - an input vector to write,
/// - the mixed destination sizes for the output tensor,
/// - and the vector sizes used for vectorization (i.e., the leading N dims,
/// for some value of N),
///
/// this function generates the following sequence of ops:
///
/// %dest = tensor.empty(%destSizes)
/// %res = vector.transfer_write %input into %dest
/// Generates the following operation:
/// %res = vector.transfer_write %vectorToStore into %dest
///
/// If the leading N dimensions of the destination tensor do not match
/// `inputVecSizesForLeadingDims` (where N =
/// rank(`inputVecSizesForLeadingDims`)), masking is applied to ensure
/// correctness:
/// `inputVecSizesForLeadingDims`, where=
/// * N = rank(`inputVecSizesForLeadingDims`)),
Copy link
Contributor

@Max191 Max191 May 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like there may be some typo here. Did you mean to say where N = rank(...?

Copy link
Contributor Author

@banach-space banach-space May 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out 🙏🏻 Fixed :)

/// masking is applied to ensure correctness:
///
/// %dest = tensor.empty(%destSizes)
/// %write = vector.transfer_write %input into %dest
/// %mask = vector.create_mask(%destSizes)
/// %write = vector.transfer_write %vectorToStore into %dest
/// %mask = vector.create_mask(%destShape)
/// %res = vector.mask %mask { %write }
Copy link
Contributor

@dcaballe dcaballe May 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wait, this is not a correct use of vector.mask. The masked operation should be nested within the vector.mask operation. I saw a similar example in the CSE discussion so let me take a look at the verifier...

Copy link
Contributor Author

@banach-space banach-space May 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, thanks! Fixed :)

///
/// If `useInBoundsInsteadOfMasking` is set to `true`, the `in_bounds` attribute
/// is used instead of masking:
///
/// %dest = tensor.empty(%destSizes)
/// %write = vector.transfer_write %vectorToStore into %dest
/// in_bounds_flags = (...)
/// %res = vector.transfer_write %input into %dest
/// {in_bounds = in_bounds_flags}
///
/// NOTE: all write offsets are set to 0.
/// NOTE: All write offsets are set to 0.
/// TODO: Allow specyfying write offsets.
/// NOTE: When N < rank(input), the missing vector sizes are effectively
/// extracted from the trailing sizes of `destSizes`. This means those sizes
/// must be static. Supporting dynamic sizes will require the user to specify
/// the remaining vector sizes. This is left as a TODO.
/// must be static.
/// TODO: Support cases where an arbitrary dim is dynamic - this will require
/// specifying all the vector sizes.
static Operation *
createWriteOrMaskedWrite(OpBuilder &builder, Location loc, Value input,
SmallVector<OpFoldResult> destSizes,
createWriteOrMaskedWrite(OpBuilder &builder, Location loc, Value vectorToStore,
Value dest,
ArrayRef<int64_t> inputVecSizesForLeadingDims,
bool useInBoundsInsteadOfMasking = false) {

auto inputType = cast<VectorType>(input.getType());
assert(inputType.getRank() == static_cast<int64_t>(destSizes.size()) &&
ShapedType destType = cast<ShapedType>(dest.getType());
assert(cast<VectorType>(vectorToStore.getType()).getRank() ==
static_cast<int64_t>(destType.getRank()) &&
"Rank mismatch!");

Value dest = builder.create<tensor::EmptyOp>(loc, destSizes,
inputType.getElementType());
int64_t rank = cast<ShapedType>(dest.getType()).getRank();
auto zero = builder.create<arith::ConstantIndexOp>(loc, 0);
auto destShape = cast<ShapedType>(dest.getType()).getShape();

// Compute the in_bounds attribute
SmallVector<bool> inBoundsVal(rank, true);
if (useInBoundsInsteadOfMasking) {
// In this case, assume that all the required vector sizes have been
// provided.
assert(inputVecSizesForLeadingDims.size() == destSizes.size() &&
assert(inputVecSizesForLeadingDims.size() ==
static_cast<size_t>(destType.getRank()) &&
"Insufficient number of input vector sizes!");
// Update the inBounds attribute.
for (unsigned i = 0; i < rank; i++)
inBoundsVal[i] = (destShape[i] == inputVecSizesForLeadingDims[i]) &&
!ShapedType::isDynamic(destShape[i]);
}

// Generate the xfer_write Op
auto zero = builder.create<arith::ConstantIndexOp>(loc, 0);
Operation *write = builder.create<vector::TransferWriteOp>(
loc,
/*vector=*/input,
/*vector=*/vectorToStore,
/*source=*/dest,
/*indices=*/SmallVector<Value>(rank, zero),
/*inBounds=*/inBoundsVal);
assert(llvm::none_of(
destShape.drop_front(inputVecSizesForLeadingDims.size()),
[](int64_t size) { return size == ShapedType::kDynamic; }) &&
"Only dims aligned with inputVecSizesForLeadingDims may be dynamic");

// If masking is disabled, exit.
if (useInBoundsInsteadOfMasking)
return write;

// Check if masking is needed.
bool needMaskForWrite =
!llvm::equal(inputVecSizesForLeadingDims,
destShape.take_front(inputVecSizesForLeadingDims.size()));

// If masking is needed, generate the mask and mask the operation.
if (needMaskForWrite) {
SmallVector<int64_t> writeMaskShape;
writeMaskShape.append(inputVecSizesForLeadingDims.begin(),
Expand All @@ -1592,10 +1594,11 @@ createWriteOrMaskedWrite(OpBuilder &builder, Location loc, Value input,
inputVecSizesForLeadingDims.size(),
destShape.end());
auto writeMaskType = VectorType::get(writeMaskShape, builder.getI1Type());
Value maskForWrite =
builder.create<vector::CreateMaskOp>(loc, writeMaskType, destSizes);
Value maskForWrite = builder.create<vector::CreateMaskOp>(
loc, writeMaskType, tensor::getMixedSizes(builder, loc, dest));
write = mlir::vector::maskOperation(builder, write, maskForWrite);
}

return write;
}

Expand Down Expand Up @@ -1693,9 +1696,11 @@ vectorizeAsTensorPackOp(RewriterBase &rewriter, linalg::PackOp packOp,
loc, shapeCastOp.getResult(), destPermutation);

// Create TransferWriteOp.
Value dest = rewriter.create<tensor::EmptyOp>(
loc, reifiedReturnShapes[0],
transposeOp.getResult().getType().getElementType());
Operation *write =
createWriteOrMaskedWrite(rewriter, loc, transposeOp.getResult(),
/*destSizes=*/reifiedReturnShapes[0],
createWriteOrMaskedWrite(rewriter, loc, transposeOp.getResult(), dest,
/*inputVecSizesForLeadingDims=*/inputVectorSizes,
/*useInBoundsInsteadOfMasking=*/false);
newResults.push_back(write->getResult(0));
Expand Down Expand Up @@ -1830,10 +1835,13 @@ vectorizeAsTensorUnpackOp(RewriterBase &rewriter, linalg::UnPackOp unpackOp,
unpackOp.getDestType().hasStaticShape()
? vectorSizes
: shapeCastOp.getResultVectorType().getShape());
Operation *write = createWriteOrMaskedWrite(
rewriter, loc, shapeCastOp.getResult(), /*destSizes=*/reifiedRetShapes[0],
/*inputVecSizesForLeadingDims=*/writeVectorSizes,
useInBoundsInsteadOfMasking);
Value dest = rewriter.create<tensor::EmptyOp>(
loc, reifiedRetShapes[0],
shapeCastOp.getResult().getType().getElementType());
Operation *write =
createWriteOrMaskedWrite(rewriter, loc, shapeCastOp.getResult(), dest,
/*inputVecSizesForLeadingDims=*/writeVectorSizes,
useInBoundsInsteadOfMasking);
newResults.push_back(write->getResult(0));
return success();
}
Expand Down Expand Up @@ -1861,10 +1869,14 @@ vectorizeAsTensorPadOp(RewriterBase &rewriter, tensor::PadOp padOp,
auto maskedRead = vector::createReadOrMaskedRead(
rewriter, loc, padOp.getSource(), inputVectorSizes, padValue,
/*useInBoundsInsteadOfMasking=*/false);
Operation *write = createWriteOrMaskedWrite(
rewriter, loc, maskedRead, reifiedReturnShapes[0],
/*inputVecSizesForLeadingDims=*/inputVectorSizes,
/*useInBoundsInsteadOfMasking=*/false);

// Create Xfer write Op
Value dest = rewriter.create<tensor::EmptyOp>(
loc, reifiedReturnShapes[0], padOp.getResultType().getElementType());
Operation *write =
createWriteOrMaskedWrite(rewriter, loc, maskedRead, dest,
/*inputVecSizesForLeadingDims=*/inputVectorSizes,
/*useInBoundsInsteadOfMasking=*/false);
newResults.push_back(write->getResult(0));
return success();
}
Expand Down
8 changes: 6 additions & 2 deletions mlir/test/Dialect/Linalg/vectorization.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -641,7 +641,9 @@ func.func @test_masked_vectorize_dynamic_pad(
// CHECK-SAME: } : vector<2x4xi1> -> vector<2x4xf32>
// CHECK-DAG: %[[empty:.*]] = tensor.empty(%[[res_d0]], %[[res_d1]]) : tensor<?x?xf32>
// CHECK-DAG: %[[c0_3:.*]] = arith.constant 0 : index
// CHECK: %[[mask_2:.*]] = vector.create_mask %[[res_d0]], %[[res_d1]] : vector<2x4xi1>
// CHECK-DAG: %[[d2:.*]] = tensor.dim %[[empty]], {{.*}} : tensor<?x?xf32>
// CHECK-DAG: %[[d3:.*]] = tensor.dim %[[empty]], {{.*}} : tensor<?x?xf32>
// CHECK: %[[mask_2:.*]] = vector.create_mask %[[d2]], %[[d3]] : vector<2x4xi1>
// CHECK: %[[masked_write:.*]] = vector.mask %[[mask_2]] {
// CHECK-SAME: vector.transfer_write %[[masked_read]], %[[empty]][%[[c0_3]], %[[c0_3]]]
// CHECK-SAME: {in_bounds = [true, true]} : vector<2x4xf32>, tensor<?x?xf32>
Expand Down Expand Up @@ -800,7 +802,9 @@ func.func @test_vectorize_dynamic_pack(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?
// CHECK-DAG: %[[c16:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[empty:.*]] = tensor.empty(%[[d0]], %[[d1]]) : tensor<?x?x16x2xf32>
// CHECK: %[[mask_0:.*]] = vector.create_mask %[[d0]], %[[d1]], %[[c16]], %[[c2]] : vector<4x1x16x2xi1>
// CHECK-DAG: %[[d2:.*]] = tensor.dim %[[empty]], {{.*}} : tensor<?x?x16x2xf32>
// CHECK-DAG: %[[d3:.*]] = tensor.dim %[[empty]], {{.*}} : tensor<?x?x16x2xf32>
// CHECK: %[[mask_0:.*]] = vector.create_mask %[[d2]], %[[d3]], %[[c16]], %[[c2]] : vector<4x1x16x2xi1>
// CHECK: %[[masked_write:.*]] = vector.mask %[[mask_0]] {
// CHECK-SAME: vector.transfer_write %[[transpose]], %[[empty]][%[[c0_2]], %[[c0_2]], %[[c0_2]], %[[c0_2]]]
Copy link
Contributor

@dcaballe dcaballe May 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tests are correct so it looks like the issue is just in the doc.

// CHECK-SAME: {in_bounds = [true, true, true, true]} : vector<4x1x16x2xf32>, tensor<?x?x16x2xf32>
Expand Down