[mlir][linalg] Split GenericPadOpVectorizationPattern into two patterns #111349
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
banach-space
requested review from
ftynse,
nicolasvasilache,
dcaballe,
hanhanW and
rengolin
as code owners
|
@llvm/pr-subscribers-mlir-linalg @llvm/pr-subscribers-mlir Author: Andrzej Warzyński (banach-space) ChangesAt the moment,
This patch splits
With this change, we gain the following:
Note for downstream users:If you were using Finer implementation details:
1.1 The only functional change is that the vectorization of 1.2 Following-on from the above,
Patch is 22.63 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/111349.diff 6 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 0693e31b4f70af..a84de1a6a8c487 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1503,18 +1503,13 @@ using OptimizeCopyFn =
/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
/// InsertSliceOp. For now, only constant padding values are supported.
-/// `OptimizeCopyFn` can be used to customize copying step optimization.
struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
- GeneralizePadOpPattern(MLIRContext *context,
- OptimizeCopyFn optimizeCopyFn = nullptr,
- PatternBenefit benefit = 1)
- : OpRewritePattern<tensor::PadOp>(context, benefit),
- optimizeCopyFn(std::move(optimizeCopyFn)) {}
+ GeneralizePadOpPattern(MLIRContext *context, PatternBenefit benefit = 1)
+ : OpRewritePattern<tensor::PadOp>(context, benefit) {}
LogicalResult matchAndRewrite(tensor::PadOp padOp,
PatternRewriter &rewriter) const override;
protected:
- OptimizeCopyFn optimizeCopyFn;
Value createFillOrGenerateOp(RewriterBase &rewriter, tensor::PadOp padOp,
Value dest,
const SmallVector<Value> &dynSizes) const;
@@ -1663,6 +1658,11 @@ void populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
/// \see rewriteInIm2Col for more details.
void populateConvertConv2DToImg2ColPatterns(RewritePatternSet &patterns);
+/// Populates `patterns` with vectorisation patterns for tensor.insert_slice.
+/// TODO: Avoid having a dedicated `populate{}` for one pattern. Instead, either
+/// expand or merge with other `populate{}`.
+void populateInsertSliceVectorizationPatterns(RewritePatternSet &patterns);
+
/// Populates `patterns` with patterns that vectorize tensor.pad.
/// These patterns are meant to apply in a complementary fashion. Benefits
/// are used to encode a certain ordering of pattern application. To avoid
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 8e7621754f76bf..157a084fec319e 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -3477,6 +3477,9 @@ transform::VectorizeChildrenAndApplyPatternsOp::applyToOne(
patterns.add<CopyVectorizationPattern>(ctx);
+ // Add misc. vectorization patterns (e.g. for tensor.insert_slice)
+ linalg::populateInsertSliceVectorizationPatterns(patterns);
+
if (getVectorizePadding())
linalg::populatePadOpVectorizationPatterns(patterns);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 0fe096863d7b01..da5233049aaf69 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -973,12 +973,7 @@ GeneralizePadOpPattern::matchAndRewrite(tensor::PadOp padOp,
padOp.getLoc(), staticSizes, resultType.getElementType(), dynSizes);
Value fill = createFillOrGenerateOp(rewriter, padOp, emptyTensor, dynSizes);
- // Try optimize the copy of source.
- if (optimizeCopyFn && optimizeCopyFn(rewriter, padOp, fill).succeeded())
- return success();
-
- // tensor::PadOps cannot be optimized. Generate a InsertSliceOp instead
- // for copying the PadOp source.
+ // Generate a InsertSliceOp for copying the PadOp source.
auto sourceType = padOp.getSourceType();
// Compute size of source of tensor::PadOp.
SmallVector<OpFoldResult> srcSizes =
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 09c6b2683b4388..9c9bf38e14300a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2262,115 +2262,6 @@ LogicalResult mlir::linalg::vectorizeCopy(RewriterBase &rewriter,
//----------------------------------------------------------------------------//
// Misc. vectorization patterns.
//----------------------------------------------------------------------------//
-
-/// Helper function that retrieves the value of an IntegerAttr.
-static int64_t getIntFromAttr(Attribute attr) {
- return cast<IntegerAttr>(attr).getInt();
-}
-
-/// Given an ArrayRef of OpFoldResults, return a vector of Values.
-/// IntegerAttrs are converted to ConstantIndexOps. Other attribute types are
-/// not supported.
-static SmallVector<Value> ofrToIndexValues(RewriterBase &rewriter, Location loc,
- ArrayRef<OpFoldResult> ofrs) {
- SmallVector<Value> result;
- for (auto o : ofrs) {
- if (auto val = llvm::dyn_cast_if_present<Value>(o)) {
- result.push_back(val);
- } else {
- result.push_back(rewriter.create<arith::ConstantIndexOp>(
- loc, getIntFromAttr(o.template get<Attribute>())));
- }
- }
- return result;
-}
-
-/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
-/// InsertSliceOp. For now, only constant padding values are supported.
-/// If there is enough static type information, TransferReadOps and
-/// TransferWriteOps may be generated instead of InsertSliceOps.
-struct GenericPadOpVectorizationPattern : public GeneralizePadOpPattern {
- GenericPadOpVectorizationPattern(MLIRContext *context,
- PatternBenefit benefit = 1)
- : GeneralizePadOpPattern(context, tryVectorizeCopy, benefit) {}
- /// Vectorize the copying of a tensor::PadOp's source. This is possible if
- /// each dimension size is statically know in the source type or the result
- /// type (or both).
- static LogicalResult tryVectorizeCopy(RewriterBase &rewriter,
- tensor::PadOp padOp, Value dest) {
- auto sourceType = padOp.getSourceType();
- auto resultType = padOp.getResultType();
- if (!VectorType::isValidElementType(sourceType.getElementType()))
- return failure();
-
- // Copy cannot be vectorized if pad value is non-constant and source shape
- // is dynamic. In case of a dynamic source shape, padding must be appended
- // by TransferReadOp, but TransferReadOp supports only constant padding.
- auto padValue = padOp.getConstantPaddingValue();
- if (!padValue) {
- if (!sourceType.hasStaticShape())
- return failure();
- // Create dummy padding value.
- auto elemType = sourceType.getElementType();
- padValue = rewriter.create<arith::ConstantOp>(
- padOp.getLoc(), elemType, rewriter.getZeroAttr(elemType));
- }
-
- SmallVector<int64_t> vecShape;
- SmallVector<bool> readInBounds;
- SmallVector<bool> writeInBounds;
- for (unsigned i = 0; i < sourceType.getRank(); ++i) {
- if (!sourceType.isDynamicDim(i)) {
- vecShape.push_back(sourceType.getDimSize(i));
- // Source shape is statically known: Neither read nor write are
- // out-of- bounds.
- readInBounds.push_back(true);
- writeInBounds.push_back(true);
- } else if (!resultType.isDynamicDim(i)) {
- // Source shape is not statically known, but result shape is.
- // Vectorize with size of result shape. This may be larger than the
- // source size.
- vecShape.push_back(resultType.getDimSize(i));
- // Read may be out-of-bounds because the result size could be larger
- // than the source size.
- readInBounds.push_back(false);
- // Write is out-of-bounds if low padding > 0.
- writeInBounds.push_back(
- getConstantIntValue(padOp.getMixedLowPad()[i]) ==
- static_cast<int64_t>(0));
- } else {
- // Neither source nor result dim of padOp is static. Cannot vectorize
- // the copy.
- return failure();
- }
- }
- auto vecType = VectorType::get(vecShape, sourceType.getElementType());
-
- // Generate TransferReadOp.
- SmallVector<Value> readIndices(
- vecType.getRank(),
- rewriter.create<arith::ConstantIndexOp>(padOp.getLoc(), 0));
- auto read = rewriter.create<vector::TransferReadOp>(
- padOp.getLoc(), vecType, padOp.getSource(), readIndices, padValue,
- ArrayRef<bool>{readInBounds});
-
- // If `dest` is a FillOp and the TransferWriteOp would overwrite the
- // entire tensor, write directly to the FillOp's operand.
- if (llvm::equal(vecShape, resultType.getShape()) &&
- llvm::all_of(writeInBounds, [](bool b) { return b; }))
- if (auto fill = dest.getDefiningOp<FillOp>())
- dest = fill.output();
-
- // Generate TransferWriteOp.
- auto writeIndices =
- ofrToIndexValues(rewriter, padOp.getLoc(), padOp.getMixedLowPad());
- rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
- padOp, read, dest, writeIndices, ArrayRef<bool>{writeInBounds});
-
- return success();
- }
-};
-
/// Base pattern for rewriting tensor::PadOps whose result is consumed by a
/// given operation type OpTy.
template <typename OpTy>
@@ -2604,6 +2495,177 @@ struct PadOpVectorizationWithTransferWritePattern
}
};
+/// Given an ArrayRef of OpFoldResults, return a vector of Values.
+/// IntegerAttrs are converted to ConstantIndexOps. Other attribute types are
+/// not supported.
+static SmallVector<Value> ofrToIndexValues(RewriterBase &rewriter, Location loc,
+ ArrayRef<OpFoldResult> ofrs) {
+ SmallVector<Value> result;
+ for (auto o : ofrs) {
+ if (auto val = llvm::dyn_cast_if_present<Value>(o)) {
+ result.push_back(val);
+ } else {
+ result.push_back(rewriter.create<arith::ConstantIndexOp>(
+ loc, cast<IntegerAttr>(cast<Attribute>(o)).getInt()));
+ }
+ }
+ return result;
+}
+
+/// Returns the effective Pad value for the input op, provided it's a scalar.
+///
+/// Many Ops exhibit pad-like behaviour, but this isn't always explicit. If
+/// this Op performs padding, retrieve the padding value provided that it's
+/// a scalar and static/fixed for all the padded values. Returns an empty value
+/// otherwise.
+static Value getStaticPadVl(Operation *op) {
+ if (!op)
+ return {};
+
+ // 1. vector.broadcast - return the value that's being broadcast,
+ // provided that it's a scalar.
+ if (auto bcast = llvm::dyn_cast<vector::BroadcastOp>(op)) {
+ auto source = bcast.getSource();
+ if (llvm::dyn_cast<VectorType>(source.getType()))
+ return {};
+
+ return source;
+ }
+
+ // 1. linalg.fill - use the scalar input value that used to fill the output
+ // tensor.
+ if (auto fill = llvm::dyn_cast<linalg::FillOp>(op)) {
+ return fill.getInputs()[0];
+ }
+
+ // 2. tensor.generateOp - can't guarantee the value is fixed without
+ // analysing, bail out.
+ if (auto generate = llvm::dyn_cast<tensor::GenerateOp>(op)) {
+ return {};
+ }
+
+ // 3. vector.transfer_write - inspect the input vector that's written from. If
+ // if contains a single value that has been broadcast (e.g. via
+ // vector.broadcast), extract it, fail otherwise.
+ if (auto xferWrite = llvm::dyn_cast<vector::TransferWriteOp>(op))
+ return getStaticPadVl(xferWrite.getVector().getDefiningOp());
+
+ // 4. tensor.insert_slice - inspect the destination tensor. If it's larger
+ // than the input tensor, then, provided it's constant, we'll extract the
+ // value that was used to generate it (via e.g. linalg.fill), fail otherwise.
+ // TODO: Clarify the semantics when the input tensor is larger than the
+ // destination.
+ if (auto slice = llvm::dyn_cast<tensor::InsertSliceOp>(op))
+ return getStaticPadVl(slice.getDest().getDefiningOp());
+
+ return {};
+}
+
+/// Rewrite tensor.insert.slice as a vector.transfer_read +
+/// vector.transfer_write pair. The vector size is inferred from the static
+/// dims in the input and output tensors. If a dim is dynamic in both the input
+/// and output tensors, bails out.
+///
+/// Before:
+/// !t_in_type = tensor<1x2x3xf32>
+/// !t_out_type = tensor<9x8x7x1x2x3xf32>
+/// !v_type = vector<1x2x3xf32>
+/// %inserted_slice = tensor.insert_slice %src into %dest ... : !t_in_type
+/// into !t_out_type
+/// After:
+/// %read = vector.transfer_read %src[...], %pad ... : !t_in_type, !v_type
+/// %write = vector.transfer_write %read, %dest ... : !v_type, !t_out_type
+///
+/// TODO: Support masking
+struct InsertSliceVectorizePattern
+ : public OpRewritePattern<tensor::InsertSliceOp> {
+ using OpRewritePattern<tensor::InsertSliceOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tensor::InsertSliceOp sliceOp,
+ PatternRewriter &rewriter) const final {
+ auto sourceType = sliceOp.getSource().getType();
+ if (!VectorType::isValidElementType(sourceType.getElementType()))
+ return failure();
+
+ auto resultType = sliceOp.getResultType();
+
+ // 1. Get the pad value.
+ // TransferReadOp requires a scalar padding value. Note that:
+ // * for in-bounds access, the value is actually irrelevant.
+ // There are 2 cases in which xfer.read accesses are known to be in-bounds:
+ // 1. The source shape is static (output vector sizes would be based on
+ // the source shape and hence all memory accesses would be in-bounds),
+ // 2. Masking is used (output vector sizes would be user-provided, in which
+ // case it is assumed that all memory accesses are in-bounds). This
+ // remains a TODO.
+ //
+ // When the value is not known and not needed, use 0. Otherwise, bail out.
+ Value padValue = getStaticPadVl(sliceOp);
+ bool isOutOfBoundsRead = !sourceType.hasStaticShape();
+
+ if (!padValue && isOutOfBoundsRead) {
+ LDBG("Failed to get a pad value for out-of-bounds read access\n");
+ return failure();
+ }
+
+ if (!padValue) {
+ auto elemType = sourceType.getElementType();
+ padValue = rewriter.create<arith::ConstantOp>(
+ sliceOp.getLoc(), elemType, rewriter.getZeroAttr(elemType));
+ }
+
+ // 2. Get the vector shape and in-bounds attributes
+ SmallVector<int64_t> vecShape;
+ SmallVector<bool> readInBounds;
+ SmallVector<bool> writeInBounds;
+ for (unsigned i = 0; i < sourceType.getRank(); ++i) {
+ if (!sourceType.isDynamicDim(i)) {
+ vecShape.push_back(sourceType.getDimSize(i));
+ // Source shape is statically known: Neither read nor write are
+ // out-of-bounds.
+ readInBounds.push_back(true);
+ writeInBounds.push_back(true);
+ } else if (!resultType.isDynamicDim(i)) {
+ // Source shape is not statically known, but result shape is.
+ // Vectorize with size of result shape. This may be larger than the
+ // source size.
+ vecShape.push_back(resultType.getDimSize(i));
+ // Read may be out-of-bounds because the result size could be larger
+ // than the source size.
+ readInBounds.push_back(false);
+ // Write will in-bounds provided that the corresponding write idx is 0.
+ // To keep this logic simple, conservatively mark as out-of-bounds.
+ writeInBounds.push_back(false);
+ } else {
+ // Neither source nor result dim of padOp is static. Cannot vectorize
+ // the copy.
+ // TODO: Add support for masking
+ return failure();
+ }
+ }
+ auto vecType = VectorType::get(vecShape, sourceType.getElementType());
+
+ // 3. Generate TransferReadOp.
+ SmallVector<Value> readIndices(
+ vecType.getRank(),
+ rewriter.create<arith::ConstantIndexOp>(sliceOp.getLoc(), 0));
+ auto read = rewriter.create<vector::TransferReadOp>(
+ sliceOp.getLoc(), vecType, sliceOp.getSource(), readIndices, padValue,
+ ArrayRef<bool>{readInBounds});
+
+ // 4. Generate TransferWriteOp.
+ auto writeIndices =
+ ofrToIndexValues(rewriter, sliceOp.getLoc(), sliceOp.getMixedOffsets());
+
+ // 5. Finalize
+ rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
+ sliceOp, read, sliceOp.getDest(), writeIndices,
+ ArrayRef<bool>{writeInBounds});
+
+ return success();
+ }
+};
+
/// Rewrite use of tensor::PadOp result in InsertSliceOp. E.g.:
/// ```
/// %0 = tensor.pad %src ... : tensor<?x?xf32> to tensor<17x5xf32>
@@ -2691,10 +2753,14 @@ struct PadOpVectorizationWithInsertSlicePattern
}
};
+void mlir::linalg::populateInsertSliceVectorizationPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<InsertSliceVectorizePattern>(patterns.getContext());
+}
+
void mlir::linalg::populatePadOpVectorizationPatterns(
RewritePatternSet &patterns, PatternBenefit baseBenefit) {
- patterns.add<GenericPadOpVectorizationPattern>(patterns.getContext(),
- baseBenefit);
+ patterns.add<GeneralizePadOpPattern>(patterns.getContext(), baseBenefit);
// Try these specialized patterns first before resorting to the generic one.
patterns.add<PadOpVectorizationWithTransferReadPattern,
PadOpVectorizationWithTransferWritePattern,
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index c481a723c5623c..4b5a66f8fb5b92 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -352,6 +352,15 @@ func.func @illegal_fill_tensor_with_memref_return
// -----
+func.func @illegal_fill_value_type(%arg0 : tensor<2x2xf32>, %arg1 : tensor<2xf32>) -> tensor<2x2xf32>
+{
+ // expected-error @+1 {{expected op with scalar input}}
+ %0 = linalg.fill ins(%arg1 : tensor<2xf32>) outs(%arg0 : tensor<2x2xf32>) -> tensor<2x2xf32>
+ return %0 : tensor<2x2xf32>
+}
+
+// -----
+
func.func @invalid_static_matmul(%arg0: memref<2x4xf32>, %arg1: memref<3x4xf32>, %arg2: memref<2x4xf32>) {
// expected-error @+1 {{inferred input/output operand #1 has shape's dimension #0 to be 4, but found 3}}
linalg.matmul ins(%arg0, %arg1 : memref<2x4xf32>, memref<3x4xf32>)
diff --git a/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir b/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir
index e7beb725471123..a2f1b69db2fca3 100644
--- a/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir
@@ -1090,16 +1090,20 @@ module attributes {transform.with_named_sequence} {
func.func private @make_vector() -> tensor<12x13xf32>
-// CHECK-LABEL: func @pad_and_insert_slice_dest
-// CHECK-SAME: %[[ARG0:.*]]: tensor<1x5x6xf32>
-// Check the insert slice is not rewritten if the padded result is used by the destination operand.
-// CHECK-NOT: tensor.pad
-// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1x12x13xf32>
-// CHECK: %[[WRITE_1:.*]] = vector.transfer_write %{{.*}}, %[[EMPTY]]{{.*}} : vector<1x12x13xf32>, tensor<1x12x13xf32>
-// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0:.*]]{{.*}} : tensor<1x5x6xf32>, vector<1x5x6xf32>
-// CHECK: %[[WRITE_2:.*]] = vector.transfer_write %[[READ]], %[[WRITE_1]]{{.*}} : vector<1x5x6xf32>, tensor<1x12x13xf32>
-// CHECK: %[[T1:.*]] = call @make_vector() : () -> tensor<12x13xf32>
-// CHECK: tensor.insert_slice %[[T1]] into %[[WRITE_2]]
+// CHECK-LABEL: func.func @pad_and_insert_slice_dest(
+// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x5x6xf32>) -> tensor<1x12x13xf32> {
+// CHECK: %[[C0:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[CST:.*]] = arith.constant dense<5.000000e+00> : vector<1x12x13xf32>
+// CHECK: %[[C0_IDX:.*]] = arith.constant 0 : index
+// CHECK: %[[PAD_VAL:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1x12x13xf32>
+// CHECK: %[[WRITE_1:.*]] = vector.transfer_write %[[CST]], %[[EMPTY]]{{\[}}%[[C0_IDX]], %[[C0_IDX]], %[[C0_IDX]]] {in_bounds = [true, true, true]} : vector<1x12x13xf32>, tensor<1x12x13xf32>
+// CHECK: %[[READ_1:.*]] = vector.transfer_read %[[ARG_0]]{{\[}}%[[C0_IDX]], %[[C0_IDX]], %[[C0_IDX]]], %[[PAD_VAL]] {in_bounds = [true, true, true]} : tensor<1x5x6xf32>, vector<1x5x6xf32>
+// CHECK: %[[WRITE_2:.*]] = vector.transfer_write %[[READ_1]], %[[WRITE_1]]{{\[}}%[[C0_IDX]],...
[truncated]
|
hanhanW
requested changes
Oct 25, 2024
| /// Given an ArrayRef of OpFoldResults, return a vector of Values. | ||
| /// IntegerAttrs are converted to ConstantIndexOps. Other attribute types are | ||
| /// not supported. | ||
| static SmallVector<Value> ofrToIndexValues(RewriterBase &rewriter, Location loc, |
There was a problem hiding this comment.
Here is the method (i.e., getValueOrCreateConstantIndexOp) that you're looking for. :)
llvm-project/mlir/include/mlir/Dialect/Arith/Utils/Utils.h
Lines 68 to 72 in 38caf28
Comment on lines
+2529
to
+2532
| if (llvm::dyn_cast<VectorType>(source.getType())) | ||
| return {}; |
There was a problem hiding this comment.
ah, I was confused. Then I realized that we can broadcast a scalar to vector (e.g., f32 -> vector<...xf32>).
There was a problem hiding this comment.
Added a small note to clarify.
| /// this Op performs padding, retrieve the padding value provided that it's | ||
| /// a scalar and static/fixed for all the padded values. Returns an empty value | ||
| /// otherwise. | ||
| static Value getStaticPadVl(Operation *op) { |
There was a problem hiding this comment.
[optional]: Can we replace th Vl with Val, which looks clearer to me? I also search the Vl under mlir/ and I found that we seldom use Vl in the codebase.
Suggested change
| static Value getStaticPadVl(Operation *op) { | |
| static Value getStaticPadVal(Operation *op) { |
There was a problem hiding this comment.
Oh, that was actually a typo (that I managed to use consistently throughout my patch 😅 ).
| /// %write = vector.transfer_write %read, %dest ... : !v_type, !t_out_type | ||
| /// | ||
| /// TODO: Support masking | ||
| struct InsertSliceVectorizePattern |
There was a problem hiding this comment.
It can be done in a follow-up. I think we want to have an unified API and vectorization path, which helps us manage the "future" extensibility feature/design better. So it'd be a plus if we move the implementation to vectorize() method.
(There was a long and old discussion in https://reviews.llvm.org/D150495 though it's not happening...but I think the point is still hold...)
There was a problem hiding this comment.
So it'd be a plus if we move the implementation to vectorize() method.
I am wondering whether that would be desirable in this case ... Let me explain.
-
Today, vectorization of
tensor.insert_sliceis shoehorned intoGenericPadOpVectorizationPattern. This means thattensor.insert_sliceis only vectorized when generated by the "decomposition" part ofGenericPadOpVectorizationPattern. -
This PR generalizes that by adding an opt-in pattern,
InsertSliceVectorizePattern, that will vectorize any qualifyingtensor.insert_slice(but users need to opt-in to use this pattern). -
Integrating this logic into
vectorize()would mean that the vectorization oftensor.insert_slicewould become non-optional.
My only concern is that vectorizing tensor.insert_slice might break various assumptions made elsewhere. Nothing concrete comes to my mind, but I'd rather be conservative and start with step 2. (going by tests in MLIR, all seems fine). If all is well, I will send a follow-on PR implementing step 3.
Btw, thank you for reminding me of that discussion on the overall design 🙏🏻
There was a problem hiding this comment.
My only concern is that vectorizing tensor.insert_slice might break various assumptions made elsewhere. Nothing concrete comes to my mind, but I'd rather be conservative and start with step 2.
Yes, good point! The plan sounds good to me, and yes we should be careful for the insert_slice case. Because there are some transformations rely on these tiling artifacts (i.e., extract_slice/insert_slice pair). I think we'll move to better state with this PR. IIRC, we'll be able to remove the other three patterns (see below) with proper pad op vectorization support. (I think it will become transfer_read, can get folded into other transfer ops in folders/canonicalization.)
// Try these specialized patterns first before resorting to the generic one.
patterns.add<PadOpVectorizationWithTransferReadPattern,
PadOpVectorizationWithTransferWritePattern,
PadOpVectorizationWithInsertSlicePattern>(
patterns.getContext(), baseBenefit.getBenefit() + 1);
At the moment, `GenericPadOpVectorizationPattern` implements two
orthogonal transformations:
1. Rewrites `tensor::PadOp` into a sequence of `tensor::EmptyOp`,
`linalg::FillOp` and `tensor::InsertSliceOp`.
2. Vectorizes (where possible) `tensor::InsertSliceOp` (see
`tryVectorizeCopy`).
This patch splits `GenericPadOpVectorizationPattern` into two separate
patterns:
1. `GeneralizePadOpPattern` for the first transformation (note that
currently `GenericPadOpVectorizationPattern` inherits from
`GeneralizePadOpPattern`).
2. `InsertSliceVectorizePattern` to vectorize `tensor::InsertSliceOp`.
With this change, we gain the following:
* a clear separation between pre-processing and vectorization
transformations/stages,
* a path to support masked vectorisation for `tensor.insert_slice`
(with a dedicated pattern for vectorization, it is much easier to
specify the input vector sizes used in masking),
* more opportunities to vectorize `tensor.insert_slice`.
Note for downstream users:
--------------------------
If you were using `populatePadOpVectorizationPatterns`, following this
change you will also have to add
`populateInsertSliceVectorizationPatterns`.
Finer implementation details:
-----------------------------
1. The majority of changes in this patch are copy & paste + some edits.
1.1 The only functional change is that the vectorization of
`tensor.insert_slice` is now broadly available (as opposed to being
constrained to the pad vectorization pattern:
`GenericPadOpVectorizationPattern`).
1.2 Following-on from the above, `@pad_and_insert_slice_dest` is
updated. As expected, the input `tensor.insert_slice` Op is no
longer "preserved" and instead gets vectorized successfully.
2. The `linalg.fill` case in `getConstantPadVal` works under the
assumption that only _scalar_ source values can be used. That's
consistent with the definition of the Op, but it's not tested at the
moment. Hence a test case in Linalg/invalid.mlir is added.
3. The behaviour of the two TD vectorization Ops,
`transform.structured.vectorize_children_and_apply_patterns` and
`transform.structured.vectorize` is preserved.
… patterns * Incorporate suggestions from Hanhan * Add a negative test to document when vectorization of tensor.insert_slice might fail * Update `@pad_and_insert_slice_dest` that was added in llvm#112504 (this change means that _all_ qualifying `tensor.insert_slice` Ops are vectorized). * Added more tests to demonstrate other cases (e.g. default vs non-default pad value).
banach-space
force-pushed
the
andrzej/extract_pad_vec_patterns
branch
from
e7aa443 to
a8406b3
Compare
hanhanW
approved these changes
Oct 28, 2024
…nto two patterns Final tweak to tests
NoumanAmir657
pushed a commit
to NoumanAmir657/llvm-project
that referenced
this pull request
Nov 4, 2024
…ns (llvm#111349) At the moment, `GenericPadOpVectorizationPattern` implements two orthogonal transformations: 1. Rewrites `tensor::PadOp` into a sequence of `tensor::EmptyOp`, `linalg::FillOp` and `tensor::InsertSliceOp`. 2. Vectorizes (where possible) `tensor::InsertSliceOp` (see `tryVectorizeCopy`). This patch splits `GenericPadOpVectorizationPattern` into two separate patterns: 1. `GeneralizePadOpPattern` for the first transformation (note that currently `GenericPadOpVectorizationPattern` inherits from `GeneralizePadOpPattern`). 2. `InsertSliceVectorizePattern` to vectorize `tensor::InsertSliceOp`. With this change, we gain the following: * a clear separation between pre-processing and vectorization transformations/stages, * a path to support masked vectorisation for `tensor.insert_slice` (with a dedicated pattern for vectorization, it is much easier to specify the input vector sizes used in masking), * more opportunities to vectorize `tensor.insert_slice`. Note for downstream users: -------------------------- If you were using `populatePadOpVectorizationPatterns`, following this change you will also have to add `populateInsertSliceVectorizationPatterns`. Finer implementation details: ----------------------------- 1. The majority of changes in this patch are copy & paste + some edits. 1.1. The only functional change is that the vectorization of `tensor.insert_slice` is now broadly available (as opposed to being constrained to the pad vectorization pattern: `GenericPadOpVectorizationPattern`). 1.2. Following-on from the above, `@pad_and_insert_slice_dest` is updated. As expected, the input `tensor.insert_slice` Op is no longer "preserved" and instead gets vectorized successfully. 2. The `linalg.fill` case in `getConstantPadVal` works under the assumption that only _scalar_ source values can be used. That's consistent with the definition of the Op, but it's not tested at the moment. Hence a test case in Linalg/invalid.mlir is added. 3. The behaviour of the two TD vectorization Ops, `transform.structured.vectorize_children_and_apply_patterns` and `transform.structured.vectorize` is preserved.
banach-space
added a commit
to banach-space/llvm-project
that referenced
this pull request
Nov 22, 2024
…nsformation Currently, `GeneralizePadOpPattern` is grouped under `populatePadOpVectorizationPatterns`. However, as noted in llvm#111349, this transformation "decomposes" rather than "vectorizes" `tensor.pad`. As such, it functions as: * a vectorization _pre-processing_ transformation, not * a vectorization transformation itself. To clarify its purpose, this PR turns `GeneralizePadOpPattern` into a standalone transformation by: * introducing a dedicated `populateDecomposePadPatterns` method, * adding a `apply_patterns.linalg.decompose_pad` Transform Dialect Op, and * removing it from `populatePadOpVectorizationPatterns`. In addition, to better reflect its role, it is renamed as "decomposition" rather then "generalization". That's to better reflect its role. This is in line with the recent renaming of similar ops, i.e. tensor.pack/tensor.unpack Ops in llvm#116439.
banach-space
added a commit
that referenced
this pull request
Nov 26, 2024
…nsformation (#117329) Currently, `GeneralizePadOpPattern` is grouped under `populatePadOpVectorizationPatterns`. However, as noted in #111349, this transformation "decomposes" rather than "vectorizes" `tensor.pad`. As such, it functions as: * a vectorization _pre-processing_ transformation, not * a vectorization transformation itself. To clarify its purpose, this PR turns `GeneralizePadOpPattern` into a standalone transformation by: * introducing a dedicated `populateDecomposePadPatterns` method, * adding a `apply_patterns.linalg.decompose_pad` Transform Dialect Op, * removing it from `populatePadOpVectorizationPatterns`. In addition, to better reflect its role, it is renamed as "decomposition" rather then "generalization". This is in line with the recent renaming of similar ops, i.e. tensor.pack/tensor.unpack Ops in #116439.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
At the moment,
GenericPadOpVectorizationPatternimplements twoorthogonal transformations:
tensor::PadOpinto a sequence oftensor::EmptyOp,linalg::FillOpandtensor::InsertSliceOp.tensor::InsertSliceOp(seetryVectorizeCopy).This patch splits
GenericPadOpVectorizationPatterninto two separatepatterns:
GeneralizePadOpPatternfor the first transformation (note thatcurrently
GenericPadOpVectorizationPatterninherits fromGeneralizePadOpPattern).InsertSliceVectorizePatternto vectorizetensor::InsertSliceOp.With this change, we gain the following:
transformations/stages,
tensor.insert_slice(with a dedicated pattern for vectorization, it is much easier to
specify the input vector sizes used in masking),
tensor.insert_slice.Note for downstream users:
If you were using
populatePadOpVectorizationPatterns, following thischange you will also have to add
populateInsertSliceVectorizationPatterns.Finer implementation details:
The majority of changes in this patch are copy & paste + some edits.
1.1. The only functional change is that the vectorization of
tensor.insert_sliceis now broadly available (as opposed to beingconstrained to the pad vectorization pattern:
GenericPadOpVectorizationPattern).1.2. Following-on from the above,
@pad_and_insert_slice_destisupdated. As expected, the input
tensor.insert_sliceOp is nolonger "preserved" and instead gets vectorized successfully.
The
linalg.fillcase ingetConstantPadValworks under theassumption that only scalar source values can be used. That's
consistent with the definition of the Op, but it's not tested at the
moment. Hence a test case in Linalg/invalid.mlir is added.
The behaviour of the two TD vectorization Ops,
transform.structured.vectorize_children_and_apply_patternsandtransform.structured.vectorizeis preserved.