[llvm-branch-commits] [mlir] 3afbfb4 - [mlir][NFC] Move helper substWithMin into Affine utils
Author: Thomas Raoux
Date: 2021-01-15T17:13:56-08:00
New Revision: 3afbfb4145bea3796f1137c34848093c3435
URL:
https://github.com/llvm/llvm-project/commit/3afbfb4145bea3796f1137c34848093c3435
DIFF:
https://github.com/llvm/llvm-project/commit/3afbfb4145bea3796f1137c34848093c3435.diff
LOG: [mlir][NFC] Move helper substWithMin into Affine utils
This allow using this helper outside of the linalg canonicalization.
Differential Revision: https://reviews.llvm.org/D94826
Added:
Modified:
mlir/include/mlir/Dialect/Affine/Utils.h
mlir/lib/Dialect/Affine/Utils/Utils.cpp
mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
Removed:
diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h
b/mlir/include/mlir/Dialect/Affine/Utils.h
index 4f36bb800dd1..be6985dfe403 100644
--- a/mlir/include/mlir/Dialect/Affine/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Utils.h
@@ -13,6 +13,7 @@
#ifndef MLIR_DIALECT_AFFINE_UTILS_H
#define MLIR_DIALECT_AFFINE_UTILS_H
+#include "mlir/IR/AffineExpr.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
@@ -130,6 +131,15 @@
vectorizeAffineLoopNest(std::vector> &loops,
/// early if the op is already in a normalized form.
void normalizeAffineParallel(AffineParallelOp op);
+/// Traverse `e` and return an AffineExpr where all occurrences of `dim` have
+/// been replaced by either:
+/// - `min` if `positivePath` is true when we reach an occurrence of `dim`
+/// - `max` if `positivePath` is true when we reach an occurrence of `dim`
+/// `positivePath` is negated each time we hit a multiplicative or divisive
+/// binary op with a constant negative coefficient.
+AffineExpr substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
+AffineExpr max, bool positivePath = true);
+
} // namespace mlir
#endif // MLIR_DIALECT_AFFINE_UTILS_H
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 844a02f8aa21..c66e111fb9cf 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -226,3 +226,30 @@ LogicalResult mlir::hoistAffineIfOp(AffineIfOp ifOp, bool
*folded) {
return success();
}
+
+// Return the min expr after replacing the given dim.
+AffineExpr mlir::substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
+ AffineExpr max, bool positivePath) {
+ if (e == dim)
+return positivePath ? min : max;
+ if (auto bin = e.dyn_cast()) {
+AffineExpr lhs = bin.getLHS();
+AffineExpr rhs = bin.getRHS();
+if (bin.getKind() == mlir::AffineExprKind::Add)
+ return substWithMin(lhs, dim, min, max, positivePath) +
+ substWithMin(rhs, dim, min, max, positivePath);
+
+auto c1 = bin.getLHS().dyn_cast();
+auto c2 = bin.getRHS().dyn_cast();
+if (c1 && c1.getValue() < 0)
+ return getAffineBinaryOpExpr(
+ bin.getKind(), c1, substWithMin(rhs, dim, min, max, !positivePath));
+if (c2 && c2.getValue() < 0)
+ return getAffineBinaryOpExpr(
+ bin.getKind(), substWithMin(lhs, dim, min, max, !positivePath), c2);
+return getAffineBinaryOpExpr(
+bin.getKind(), substWithMin(lhs, dim, min, max, positivePath),
+substWithMin(rhs, dim, min, max, positivePath));
+ }
+ return e;
+}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 5b6302a7e5a2..b6171ff9c5b1 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -12,6 +12,7 @@
//===--===//
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
@@ -332,38 +333,6 @@ LogicalResult mlir::linalg::applyStagedPatterns(
return success();
}
-/// Traverse `e` and return an AffineExpr where all occurrences of `dim` have
-/// been replaced by either:
-/// - `min` if `positivePath` is true when we reach an occurrence of `dim`
-/// - `max` if `positivePath` is true when we reach an occurrence of `dim`
-/// `positivePath` is negated each time we hit a multiplicative or divisive
-/// binary op with a constant negative coefficient.
-static AffineExpr substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
- AffineExpr max, bool positivePath = true) {
- if (e == dim)
-return positivePath ? min : max;
- if (auto bin = e.dyn_cast()) {
-AffineExpr lhs = bin.getLHS();
-AffineExpr rhs = bin.getRHS();
-if (bin.getKind() == mlir::AffineExprKind::Add)
- return substWithMin(lhs, dim, min, max, positivePath) +
-
[llvm-branch-commits] [mlir] fd2083d - [mlir] Fixing potential build break in my previous commit
Author: Thomas Raoux Date: 2021-01-15T17:38:16-08:00 New Revision: fd2083d73c2d229e80c1091a0399f7e1076e8c52 URL: https://github.com/llvm/llvm-project/commit/fd2083d73c2d229e80c1091a0399f7e1076e8c52 DIFF: https://github.com/llvm/llvm-project/commit/fd2083d73c2d229e80c1091a0399f7e1076e8c52.diff LOG: [mlir] Fixing potential build break in my previous commit Added: Modified: mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt Removed: diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt index 42e2d4dcd244..b063996db55e 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt @@ -23,6 +23,7 @@ add_mlir_dialect_library(MLIRLinalgTransforms LINK_LIBS PUBLIC MLIRAffine + MLIRAffineUtils MLIRAnalysis MLIREDSC MLIRIR ___ llvm-branch-commits mailing list llvm-branch-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
[llvm-branch-commits] [mlir] cf21667 - [mlir][linalg] Add vectorization for linalg on tensor ops
Author: Thomas Raoux
Date: 2020-12-29T09:02:23-08:00
New Revision: cf216670a0bd1f2ce561a315e00649740f117e1c
URL:
https://github.com/llvm/llvm-project/commit/cf216670a0bd1f2ce561a315e00649740f117e1c
DIFF:
https://github.com/llvm/llvm-project/commit/cf216670a0bd1f2ce561a315e00649740f117e1c.diff
LOG: [mlir][linalg] Add vectorization for linalg on tensor ops
Support vectorization of linalg ops using tensor inputs/outputs.
Differential Revision: https://reviews.llvm.org/D93890
Added:
Modified:
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/lib/Dialect/Vector/VectorOps.cpp
mlir/test/Dialect/Linalg/vectorization.mlir
Removed:
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 23e452df9184..2a1d4cd2ef57 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -119,34 +119,38 @@ static bool isElementwise(Operation *op) {
return hasOnlyScalarElementwiseOp(genericOp.getRegion());
}
-static VectorType extractVectorTypeFromScalarView(Value v) {
- MemRefType mt = v.getType().cast();
- return mt.getShape().empty()
- ? VectorType()
- : VectorType::get(mt.getShape(), mt.getElementType());
+static VectorType extractVectorTypeFromShapedValue(Value v) {
+ auto st = v.getType().cast();
+ if (st.isa() && st.getShape().empty())
+return VectorType();
+ return VectorType::get(st.getShape(), st.getElementType());
}
-static Value transferReadVector(OpBuilder &builder, Value memref) {
+static Value transferReadVector(OpBuilder &builder, Value source) {
edsc::ScopedContext scope(builder);
- auto memrefType = memref.getType().cast();
- if (VectorType vectorType = extractVectorTypeFromScalarView(memref)) {
-SmallVector indices(memrefType.getRank(), std_constant_index(0));
-return vector_transfer_read(vectorType, memref, indices);
+ auto shapedType = source.getType().cast();
+ if (VectorType vectorType = extractVectorTypeFromShapedValue(source)) {
+SmallVector indices(shapedType.getRank(), std_constant_index(0));
+return vector_transfer_read(vectorType, source, indices);
}
- return std_load(memref);
+ return std_load(source);
}
-static void transferWriteVector(OpBuilder &builder, Value value, Value memref)
{
+static Value transferWriteVector(OpBuilder &builder, Value value, Value dest) {
edsc::ScopedContext scope(builder);
- auto memrefType = memref.getType().cast();
- if (VectorType vectorType = extractVectorTypeFromScalarView(memref)) {
-SmallVector indices(memrefType.getRank(), std_constant_index(0));
+ Operation *write;
+ auto shapedType = dest.getType().cast();
+ if (VectorType vectorType = extractVectorTypeFromShapedValue(dest)) {
+SmallVector indices(shapedType.getRank(), std_constant_index(0));
if (vectorType != value.getType())
value = vector_broadcast(vectorType, value);
-vector_transfer_write(value, memref, indices);
+write = vector_transfer_write(value, dest, indices);
} else {
-std_store(value, memref);
+write = std_store(value, dest);
}
+ if (!write->getResults().empty())
+return write->getResult(0);
+ return Value();
}
namespace {
@@ -167,10 +171,12 @@ class GenericVectorizer {
void vectorize(Operation &scalarOp) {
auto yieldOp = dyn_cast(scalarOp);
if (yieldOp) {
- for (auto outputAndMemref :
- llvm::zip(yieldOp.values(), generic.getOutputBuffers())) {
-Value vectorValue = vectorize(std::get<0>(outputAndMemref));
-transferWriteVector(builder, vectorValue,
std::get<1>(outputAndMemref));
+ for (auto outputs : llvm::enumerate(yieldOp.values())) {
+Value vectorValue = vectorize(outputs.value());
+Value result = transferWriteVector(builder, vectorValue,
+ generic.getOutput(outputs.index()));
+if (result)
+ results.push_back(result);
}
return;
}
@@ -182,6 +188,8 @@ class GenericVectorizer {
}
}
+ llvm::ArrayRef getResults() { return results; }
+
private:
// Transforms a scalar value into its vectorized counterpart, recursively
// vectorizing operations as necessary using the underlying builder.
@@ -261,6 +269,7 @@ class GenericVectorizer {
OpBuilder &builder;
linalg::GenericOp generic;
llvm::DenseMap valueCache;
+ SmallVector results;
};
} // namespace
@@ -271,6 +280,8 @@ static void vectorizeElementwise(linalg::GenericOp op,
OpBuilder &builder) {
for (Operation &scalarOp : op.region().front()) {
vectorizer.vectorize(scalarOp);
}
+ if (!op->getResults().empty())
+op->replaceAllUsesWith(vectorizer.getResults());
}
LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(Operation *op) {
@@ -331,32 +342,14 @@ void mlir::linalg::vec
[llvm-branch-commits] [mlir] f9190c8 - [mlir][vector] Support unrolling for transfer ops using tensors
Author: Thomas Raoux
Date: 2021-01-06T13:28:04-08:00
New Revision: f9190c868137dcf43833db2c8e1e00c7acca67bc
URL:
https://github.com/llvm/llvm-project/commit/f9190c868137dcf43833db2c8e1e00c7acca67bc
DIFF:
https://github.com/llvm/llvm-project/commit/f9190c868137dcf43833db2c8e1e00c7acca67bc.diff
LOG: [mlir][vector] Support unrolling for transfer ops using tensors
Differential Revision: https://reviews.llvm.org/D93904
Added:
Modified:
mlir/include/mlir/Dialect/Vector/VectorTransforms.h
mlir/lib/Dialect/Vector/VectorTransforms.cpp
mlir/test/Dialect/Vector/vector-transfer-unroll.mlir
mlir/test/Dialect/Vector/vector-transforms.mlir
mlir/test/lib/Transforms/TestVectorTransforms.cpp
Removed:
diff --git a/mlir/include/mlir/Dialect/Vector/VectorTransforms.h
b/mlir/include/mlir/Dialect/Vector/VectorTransforms.h
index c88aa7f5bc65..a258903d5a3a 100644
--- a/mlir/include/mlir/Dialect/Vector/VectorTransforms.h
+++ b/mlir/include/mlir/Dialect/Vector/VectorTransforms.h
@@ -71,7 +71,8 @@ SmallVector unrollSingleResultVectorOp(OpBuilder
&builder,
/// Unroll a transfer_write op. Break up the vector source into a tuple of
/// vectors matching the given shape. Then store each element with its own
-/// transfer_write.
+/// transfer_write. If the transfer_write takes a tensor source, return the
+/// unrolled Value in result.
///
/// Example:
/// vector.transfer_write %A, %M[%c0, %c0] : vector<4x4xf32>, memref<4x4xf32>
@@ -83,7 +84,8 @@ SmallVector unrollSingleResultVectorOp(OpBuilder
&builder,
/// %2 = vector.tuple_get %0, 1 : tuple, vector<2x4xf32>>
/// vector.transfer_write %2, %M[%c2, %c0] : vector<2x4xf32>, memref<4x4xf32>
LogicalResult unrollTransferWriteOp(OpBuilder &builder, Operation *op,
-ArrayRef targetShape);
+ArrayRef targetShape,
+SmallVector &result);
/// Options that control the vector unrolling.
struct UnrollVectorOptions {
@@ -143,9 +145,10 @@ struct UnrollVectorPattern : public RewritePattern {
llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; }))
return failure();
if (isa(op)) {
- if (failed(unrollTransferWriteOp(rewriter, op, *targetShape)))
+ SmallVector result;
+ if (failed(unrollTransferWriteOp(rewriter, op, *targetShape, result)))
return failure();
- rewriter.eraseOp(op);
+ rewriter.replaceOp(op, result);
return success();
}
if (op->getNumResults() != 1)
diff --git a/mlir/lib/Dialect/Vector/VectorTransforms.cpp
b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
index f1708db113d4..ca6e92d95ed0 100644
--- a/mlir/lib/Dialect/Vector/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
@@ -515,7 +515,7 @@ static void getVectorElementwiseOpUnrollState(Operation *op,
/// Generates slices of 'vectorType' according to 'sizes' and 'strides, and
/// calls 'fn' with linear index and indices for each slice.
static void generateTransferOpSlices(
-Type memrefElementType, VectorType vectorType, TupleType tupleType,
+Type shapedElementType, VectorType vectorType, TupleType tupleType,
ArrayRef sizes, ArrayRef strides, ArrayRef
indices,
OpBuilder &builder, function_ref)> fn) {
// Compute strides w.r.t. to slice counts in each dimension.
@@ -539,9 +539,9 @@ static void generateTransferOpSlices(
// vector rank is 4 - 2 = 2, and so 'indexOffset' = 3 - 2 = 1.
//
unsigned vectorRank = vectorType.getRank();
- if (auto memrefVectorElementType = memrefElementType.dyn_cast())
{
-assert(vectorRank >= memrefVectorElementType.getRank());
-vectorRank -= memrefVectorElementType.getRank();
+ if (auto sourceVectorElementType = shapedElementType.dyn_cast())
{
+assert(vectorRank >= sourceVectorElementType.getRank());
+vectorRank -= sourceVectorElementType.getRank();
}
unsigned indexOffset = numSliceIndices - vectorRank;
@@ -598,8 +598,8 @@ static Value unrollTransferReadOp(vector::TransferReadOp
readOp,
SmallVector strides(targetShape.size(), 1);
Location loc = readOp.getLoc();
- auto memrefElementType =
- readOp.source().getType().cast().getElementType();
+ auto shapedElementType =
+ readOp.source().getType().cast().getElementType();
auto tupleType = generateExtractSlicesOpResultType(
sourceVectorType, targetShape, strides, builder);
int64_t numSlices = tupleType.size();
@@ -618,7 +618,7 @@ static Value unrollTransferReadOp(vector::TransferReadOp
readOp,
readOp.permutation_map(), readOp.padding(),
readOp.masked() ? *readOp.masked() : ArrayAttr());
};
- generateTransferOpSlices(memrefElementType, sourceVectorType, tupleType,
+ generateTransferOpSlices(shapedElementType, sourceVectorType, tupleType,
targetShape, strides, indices, builder,
createSli
[llvm-branch-commits] [mlir] efd0504 - [mlir] Add hoisting transformation for transfer ops on tensor
Author: Thomas Raoux
Date: 2021-01-06T14:23:59-08:00
New Revision: efd05040e13e942a4fbb79eb798fb9833e319b51
URL:
https://github.com/llvm/llvm-project/commit/efd05040e13e942a4fbb79eb798fb9833e319b51
DIFF:
https://github.com/llvm/llvm-project/commit/efd05040e13e942a4fbb79eb798fb9833e319b51.diff
LOG: [mlir] Add hoisting transformation for transfer ops on tensor
Add same hoisting transformation existing for transfer ops on buffers for
transfer_ops on tensor. The logic is significantly different so this is done as
a separate transformation and it is expect that user would know which
transformation to use based on the flow.
Differential Revision: https://reviews.llvm.org/D94115
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
mlir/include/mlir/Dialect/Vector/VectorUtils.h
mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
mlir/lib/Dialect/Vector/VectorUtils.cpp
mlir/test/Dialect/Linalg/hoisting.mlir
mlir/test/lib/Transforms/TestLinalgHoisting.cpp
Removed:
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
b/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
index 32693555ff40..ed585d1f5cf5 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
@@ -21,8 +21,9 @@ namespace linalg {
// TODO: generalize on a per-need basis.
void hoistViewAllocOps(FuncOp func);
-/// Hoist vector.transfer_read/vector.transfer_write pairs out of immediately
-/// enclosing scf::ForOp iteratively, if the following conditions are true:
+/// Hoist vector.transfer_read/vector.transfer_write on buffers pairs out of
+/// immediately enclosing scf::ForOp iteratively, if the following conditions
+/// are true:
/// 1. The two ops access the same memref with the same indices.
/// 2. All operands are invariant under the enclosing scf::ForOp.
/// 3. No uses of the memref either dominate the transfer_read or are
@@ -35,6 +36,10 @@ void hoistViewAllocOps(FuncOp func);
// TODO: generalize on a per-need basis.
void hoistRedundantVectorTransfers(FuncOp func);
+/// Same behavior as `hoistRedundantVectorTransfers` but works on tensors
+/// instead of buffers.
+void hoistRedundantVectorTransfersOnTensor(FuncOp func);
+
} // namespace linalg
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/Vector/VectorUtils.h
b/mlir/include/mlir/Dialect/Vector/VectorUtils.h
index a06bc8cf6562..03250f0a 100644
--- a/mlir/include/mlir/Dialect/Vector/VectorUtils.h
+++ b/mlir/include/mlir/Dialect/Vector/VectorUtils.h
@@ -165,6 +165,12 @@ AffineMap getTransferMinorIdentityMap(ShapedType
shapedType,
bool isDisjointTransferSet(VectorTransferOpInterface transferA,
VectorTransferOpInterface transferB);
+/// Same behavior as `isDisjointTransferSet` but doesn't require the operations
+/// to have the same tensor/memref. This allows comparing operations accessing
+///
diff erent tensors.
+bool isDisjointTransferIndices(VectorTransferOpInterface transferA,
+ VectorTransferOpInterface transferB);
+
namespace matcher {
/// Matches vector.transfer_read, vector.transfer_write and ops that return a
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index a1797fde7da6..98d61fa6a8d9 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -81,12 +81,151 @@ void mlir::linalg::hoistViewAllocOps(FuncOp func) {
}
}
+/// Look for a transfer_read, in the given tensor uses, accessing the same
+/// offset as the transfer_write.
+static vector::TransferReadOp
+findMatchingTransferRead(vector::TransferWriteOp write, Value srcTensor) {
+ for (Operation *user : srcTensor.getUsers()) {
+auto read = dyn_cast(user);
+if (read && read.indices() == write.indices() &&
+read.getVectorType() == write.getVectorType()) {
+ return read;
+}
+ }
+ return nullptr;
+}
+
+/// Check if the chunk of data inserted by the transfer_write in the given
+/// tensor are read by any other op than the read candidate.
+static bool tensorChunkAccessedByUnknownOp(vector::TransferWriteOp write,
+ vector::TransferReadOp
candidateRead,
+ Value srcTensor) {
+ // Make sure none of the other uses read the part of the tensor modified
+ // by the transfer_write.
+ llvm::SmallVector uses;
+ uses.push_back(srcTensor.getUses());
+ while (!uses.empty()) {
+for (OpOperand &use : uses.pop_back_val()) {
+ Operation *user = use.getOwner();
+ // Skip the candidate use, only inspect the "other" uses.
+ if (user == candidateRead.getOperation() || user == write.getOperation())
+continue;
+ // Consider all transitive uses through a ve
[llvm-branch-commits] [mlir] 080943f - [mlir][vector] Support transfer op on tensor optimizations
Author: Thomas Raoux
Date: 2021-01-06T15:09:03-08:00
New Revision: 080943f7525f277579a000cf30364cc96fba6773
URL:
https://github.com/llvm/llvm-project/commit/080943f7525f277579a000cf30364cc96fba6773
DIFF:
https://github.com/llvm/llvm-project/commit/080943f7525f277579a000cf30364cc96fba6773.diff
LOG: [mlir][vector] Support transfer op on tensor optimizations
Support store to load forwarding and dead store transformations for transfer op
on tensor.
Differential Revision: https://reviews.llvm.org/D94148
Added:
Modified:
mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp
mlir/test/Dialect/Vector/vector-transferop-opt.mlir
Removed:
diff --git a/mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp
b/mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp
index ea1189d53b31..161d02cd3435 100644
--- a/mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp
@@ -34,13 +34,33 @@ static Operation *findAncestorOpInRegion(Region *region,
Operation *op) {
return op;
}
+/// Return true if the transfer_write fully writes the data accessed by the
+/// transfer_read.
+static bool transferEncompasses(vector::TransferWriteOp defWrite,
+vector::TransferReadOp read) {
+ return !defWrite.hasMaskedDim() && defWrite.indices() == read.indices() &&
+ defWrite.getVectorType() == read.getVectorType() &&
+ defWrite.permutation_map() == read.permutation_map();
+}
+
+/// Return true if the write op fully over-write the priorWrite transfer_write
+/// op.
+static bool transferEncompasses(vector::TransferWriteOp write,
+vector::TransferWriteOp priorWrite) {
+ return priorWrite.indices() == write.indices() &&
+ priorWrite.getVectorType() == write.getVectorType() &&
+ priorWrite.permutation_map() == write.permutation_map();
+}
+
namespace {
class TransferOptimization {
public:
TransferOptimization(FuncOp func) : dominators(func), postDominators(func) {}
void deadStoreOp(vector::TransferWriteOp);
+ void deadStoreOpTensor(vector::TransferWriteOp);
void storeToLoadForwarding(vector::TransferReadOp);
+ void storeToLoadForwardingTensor(vector::TransferReadOp);
void removeDeadOp() {
for (Operation *op : opToErase)
op->erase();
@@ -99,9 +119,7 @@ void
TransferOptimization::deadStoreOp(vector::TransferWriteOp write) {
continue;
if (auto nextWrite = dyn_cast(user)) {
// Check candidate that can override the store.
- if (write.indices() == nextWrite.indices() &&
- write.getVectorType() == nextWrite.getVectorType() &&
- write.permutation_map() == write.permutation_map() &&
+ if (transferEncompasses(nextWrite, write) &&
postDominators.postDominates(nextWrite, write)) {
if (firstOverwriteCandidate == nullptr ||
postDominators.postDominates(firstOverwriteCandidate, nextWrite))
@@ -173,10 +191,8 @@ void
TransferOptimization::storeToLoadForwarding(vector::TransferReadOp read) {
cast(write.getOperation()),
cast(read.getOperation(
continue;
- if (dominators.dominates(write, read) && !write.hasMaskedDim() &&
- write.indices() == read.indices() &&
- write.getVectorType() == read.getVectorType() &&
- write.permutation_map() == read.permutation_map()) {
+ if (dominators.dominates(write, read) &&
+ transferEncompasses(write, read)) {
if (lastwrite == nullptr || dominators.dominates(lastwrite, write))
lastwrite = write;
else
@@ -214,15 +230,62 @@ void
TransferOptimization::storeToLoadForwarding(vector::TransferReadOp read) {
opToErase.push_back(read.getOperation());
}
+/// Walk up the SSA links, if any write gets fully overwritten we can skip it.
+/// If it has no more uses it becomes dead.
+void TransferOptimization::deadStoreOpTensor(vector::TransferWriteOp write) {
+ auto defWrite = write.source().getDefiningOp();
+ while (defWrite) {
+if (transferEncompasses(write, defWrite)) {
+ write.sourceMutable().assign(defWrite.source());
+ if (defWrite->use_empty())
+opToErase.push_back(defWrite.getOperation());
+ return;
+}
+if (!isDisjointTransferIndices(
+cast(defWrite.getOperation()),
+cast(write.getOperation(
+ break;
+defWrite = defWrite.source().getDefiningOp();
+ }
+}
+
+/// Walk up the SSA links, if any write fully match the written vector we can
+/// replace the read by the vector. The read becomes dead and can be removed.
+void TransferOptimization::storeToLoadForwardingTensor(
+vector::TransferReadOp read) {
+ auto defWrite = read.source().getDefiningOp();
+ while (defWrite) {
+if (transferEncompasses(defWrite, read)) {
+ read.replaceAllUsesWith(defWrite.vector
[llvm-branch-commits] [mlir] 3d693bd - [mlir][vector] Add memory effects to transfer_read transfer_write ops
Author: Thomas Raoux
Date: 2021-01-11T09:25:37-08:00
New Revision: 3d693bd0bd77fe6f0dd922be374b7ba74739871a
URL:
https://github.com/llvm/llvm-project/commit/3d693bd0bd77fe6f0dd922be374b7ba74739871a
DIFF:
https://github.com/llvm/llvm-project/commit/3d693bd0bd77fe6f0dd922be374b7ba74739871a.diff
LOG: [mlir][vector] Add memory effects to transfer_read transfer_write ops
This allow more accurate modeling of the side effects and allow dead code
elimination to remove dead transfer ops.
Differential Revision: https://reviews.llvm.org/D94318
Added:
Modified:
mlir/include/mlir/Dialect/Vector/VectorOps.td
mlir/lib/Dialect/Vector/VectorOps.cpp
mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir
mlir/test/Dialect/Vector/canonicalize.mlir
Removed:
diff --git a/mlir/include/mlir/Dialect/Vector/VectorOps.td
b/mlir/include/mlir/Dialect/Vector/VectorOps.td
index 7f57dcd77def..6bfa89939b04 100644
--- a/mlir/include/mlir/Dialect/Vector/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/VectorOps.td
@@ -1055,7 +1055,8 @@ def Vector_ExtractStridedSliceOp :
def Vector_TransferReadOp :
Vector_Op<"transfer_read", [
DeclareOpInterfaceMethods,
- DeclareOpInterfaceMethods
+ DeclareOpInterfaceMethods,
+ DeclareOpInterfaceMethods
]>,
Arguments<(ins AnyShaped:$source, Variadic:$indices,
AffineMapAttr:$permutation_map, AnyType:$padding,
@@ -1224,7 +1225,8 @@ def Vector_TransferReadOp :
def Vector_TransferWriteOp :
Vector_Op<"transfer_write", [
DeclareOpInterfaceMethods,
- DeclareOpInterfaceMethods
+ DeclareOpInterfaceMethods,
+ DeclareOpInterfaceMethods
]>,
Arguments<(ins AnyVector:$vector, AnyShaped:$source,
Variadic:$indices,
diff --git a/mlir/lib/Dialect/Vector/VectorOps.cpp
b/mlir/lib/Dialect/Vector/VectorOps.cpp
index 91eab5027962..731ddae85ead 100644
--- a/mlir/lib/Dialect/Vector/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/VectorOps.cpp
@@ -2227,6 +2227,14 @@ Optional>
TransferReadOp::getShapeForUnroll() {
return SmallVector{s.begin(), s.end()};
}
+void TransferReadOp::getEffects(
+SmallVectorImpl>
+&effects) {
+ if (getShapedType().isa())
+effects.emplace_back(MemoryEffects::Read::get(), source(),
+ SideEffects::DefaultResource::get());
+}
+
//===--===//
// TransferWriteOp
//===--===//
@@ -2341,6 +2349,14 @@ Optional>
TransferWriteOp::getShapeForUnroll() {
return llvm::to_vector<4>(getVectorType().getShape());
}
+void TransferWriteOp::getEffects(
+SmallVectorImpl>
+&effects) {
+ if (getShapedType().isa())
+effects.emplace_back(MemoryEffects::Write::get(), source(),
+ SideEffects::DefaultResource::get());
+}
+
//===--===//
// MaskedLoadOp
//===--===//
diff --git a/mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir
b/mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir
index 76e71b1c0bfd..e5bb65aa4208 100644
--- a/mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt %s -convert-vector-to-scf -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -convert-vector-to-scf=full-unroll=true -split-input-file
| FileCheck %s --check-prefix=FULL-UNROLL
+// RUN: mlir-opt %s -convert-vector-to-scf -split-input-file
-allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt %s -convert-vector-to-scf=full-unroll=true -split-input-file
-allow-unregistered-dialect | FileCheck %s --check-prefix=FULL-UNROLL
// CHECK-LABEL: func @materialize_read_1d() {
func @materialize_read_1d() {
@@ -22,6 +22,9 @@ func @materialize_read_1d() {
// CHECK-NEXT: else
// CHECK-NEXT: vector.insertelement
// CHECK-NEXT: store
+ // Add a dummy use to prevent dead code elimination from removing
transfer
+ // read ops.
+ "dummy_use"(%f1, %f2, %f3, %f4) : (vector<4xf32>, vector<4xf32>,
vector<4xf32>, vector<4xf32>) -> ()
}
}
return
@@ -41,6 +44,9 @@ func @materialize_read_1d_partially_specialized(%dyn1 :
index, %dyn2 : index, %d
%f1 = vector.transfer_read %A[%i0, %i1, %i2, %i3, %i4], %f0
{permutation_map = affine_map<(d0, d1, d2, d3, d4) -> (d3)>} :
memref<7x?x?x42x?xf32>, vector<4xf32>
%i3p1 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i3)
%f2 = vector.transfer_read %A[%i0, %i1, %i2, %i3p1, %i4], %f0
{permutation_map = affine_map<(d0, d1, d2, d3, d4) -> (d3)>} :
memref<7x?x?x42x?xf32>, vector<4xf32>
+// Add a dummy use to prevent dead code elimination from removin
[llvm-branch-commits] [mlir] c1ae378 - [mlir][vector] Add side-effect information to different load/store ops
Author: Thomas Raoux
Date: 2021-01-11T12:34:14-08:00
New Revision: c1ae378205db72cd80a52b85b8474077d1aa5b15
URL:
https://github.com/llvm/llvm-project/commit/c1ae378205db72cd80a52b85b8474077d1aa5b15
DIFF:
https://github.com/llvm/llvm-project/commit/c1ae378205db72cd80a52b85b8474077d1aa5b15.diff
LOG: [mlir][vector] Add side-effect information to different load/store ops
Differential Revision: https://reviews.llvm.org/D94434
Added:
Modified:
mlir/include/mlir/Dialect/Vector/VectorOps.td
mlir/test/Dialect/Vector/canonicalize.mlir
Removed:
diff --git a/mlir/include/mlir/Dialect/Vector/VectorOps.td
b/mlir/include/mlir/Dialect/Vector/VectorOps.td
index 6bfa89939b04..8bc21b179037 100644
--- a/mlir/include/mlir/Dialect/Vector/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/VectorOps.td
@@ -1318,7 +1318,7 @@ def Vector_TransferWriteOp :
def Vector_MaskedLoadOp :
Vector_Op<"maskedload">,
-Arguments<(ins AnyMemRef:$base,
+Arguments<(ins Arg:$base,
Variadic:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$pass_thru)>,
@@ -1370,7 +1370,7 @@ def Vector_MaskedLoadOp :
def Vector_MaskedStoreOp :
Vector_Op<"maskedstore">,
-Arguments<(ins AnyMemRef:$base,
+Arguments<(ins Arg:$base,
Variadic:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$value)> {
@@ -1418,7 +1418,7 @@ def Vector_MaskedStoreOp :
def Vector_GatherOp :
Vector_Op<"gather">,
-Arguments<(ins AnyMemRef:$base,
+Arguments<(ins Arg:$base,
VectorOfRankAndType<[1], [AnyInteger]>:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$pass_thru)>,
@@ -1475,7 +1475,7 @@ def Vector_GatherOp :
def Vector_ScatterOp :
Vector_Op<"scatter">,
-Arguments<(ins AnyMemRef:$base,
+Arguments<(ins Arg:$base,
VectorOfRankAndType<[1], [AnyInteger]>:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$value)> {
@@ -1531,7 +1531,7 @@ def Vector_ScatterOp :
def Vector_ExpandLoadOp :
Vector_Op<"expandload">,
-Arguments<(ins AnyMemRef:$base,
+Arguments<(ins Arg:$base,
Variadic:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$pass_thru)>,
@@ -1586,7 +1586,7 @@ def Vector_ExpandLoadOp :
def Vector_CompressStoreOp :
Vector_Op<"compressstore">,
-Arguments<(ins AnyMemRef:$base,
+Arguments<(ins Arg:$base,
Variadic:$indices,
VectorOfRankAndType<[1], [I1]>:$mask,
VectorOfRank<[1]>:$value)> {
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir
b/mlir/test/Dialect/Vector/canonicalize.mlir
index cf4473f15f49..0ff85da85bcb 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -677,3 +677,22 @@ func @dead_transfer_op(%arg0 : tensor<4x4xf32>, %arg1 :
memref<4x4xf32>,
vector<1x4xf32>, tensor<4x4xf32>
return
}
+
+// -
+
+// CHECK-LABEL: func @dead_load
+// CHECK-NOT: vector.maskedload
+// CHECK-NOT: vector.gather
+// CHECK-NOT: vector.expandload
+// CHECK: return
+func @dead_load(%base: memref, %indices: vector<16xi32>,
+ %mask: vector<16xi1>, %passthru: vector<16xf32>) {
+ %c0 = constant 0 : index
+ %0 = vector.maskedload %base[%c0], %mask, %passthru :
+memref, vector<16xi1>, vector<16xf32> into vector<16xf32>
+ %1 = vector.gather %base[%indices], %mask, %passthru :
+memref, vector<16xi32>, vector<16xi1>, vector<16xf32> into
vector<16xf32>
+ %2 = vector.expandload %base[%c0], %mask, %passthru :
+memref, vector<16xi1>, vector<16xf32> into vector<16xf32>
+ return
+}
___
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
[llvm-branch-commits] [mlir] 8955e9f - [mlir][linalg] Fix bug in elementwise vectorization
Author: Thomas Raoux
Date: 2020-12-14T10:44:36-08:00
New Revision: 8955e9f6b75d436f92235531f003540401cd4b30
URL:
https://github.com/llvm/llvm-project/commit/8955e9f6b75d436f92235531f003540401cd4b30
DIFF:
https://github.com/llvm/llvm-project/commit/8955e9f6b75d436f92235531f003540401cd4b30.diff
LOG: [mlir][linalg] Fix bug in elementwise vectorization
Fix a bug causing to pick the wrong vector size to broadcast to when the source
vectors have different ranks.
Differential Revision: https://reviews.llvm.org/D93118
Added:
Modified:
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/test/Dialect/Linalg/vectorization.mlir
Removed:
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index a28b90b1d95c..2df1a9469eab 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -216,6 +216,7 @@ class GenericVectorizer {
if (!vecType)
continue;
if (maxSize < vecType.getNumElements()) {
+maxSize = vecType.getNumElements();
largestShape.assign(vecType.getShape().begin(),
vecType.getShape().end());
}
diff --git a/mlir/test/Dialect/Linalg/vectorization.mlir
b/mlir/test/Dialect/Linalg/vectorization.mlir
index 1c3533275e49..6019dde49983 100644
--- a/mlir/test/Dialect/Linalg/vectorization.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization.mlir
@@ -169,7 +169,7 @@ func @generic_vectorize(%arg0: memref<4x256xf32>, %arg1:
memref<4x256xf32>,
%11 = mulf %arg5, %8 : f32
%12 = rsqrt %arg5 : f32
%13 = select %7, %arg5, %arg6 : f32
-%14 = subf %arg5, %arg6 : f32
+%14 = subf %arg5, %arg4 : f32
%15 = tanh %arg5 : f32
linalg.yield %6, %8, %c1_f32, %9, %10, %11, %12, %13, %14, %15 : f32, f32,
f32, f32, f32, f32, f32, f32, f32, f32
@@ -196,7 +196,8 @@ func @generic_vectorize(%arg0: memref<4x256xf32>, %arg1:
memref<4x256xf32>,
// CHECK: %[[MUL:.*]] = mulf %[[V3]], %[[CST0]] : vector<4x256xf32>
// CHECK: %[[RSQRT:.*]] = rsqrt %[[V3]] : vector<4x256xf32>
// CHECK: %[[SEL:.*]] = select %[[CMP]], %[[V3]], %[[V1]] :
vector<4x256xi1>, vector<4x256xf32>
-// CHECK: %[[SUB:.*]] = subf %[[V3]], %[[V1]] : vector<4x256xf32>
+// CHECK: %[[V0B:.*]] = vector.broadcast %[[V0]] : vector<256xf32> to
vector<4x256xf32>
+// CHECK: %[[SUB:.*]] = subf %[[V3]], %[[V0B]] : vector<4x256xf32>
// CHECK: %[[TAN:.*]] = tanh %[[V3]] : vector<4x256xf32>
// CHECK: vector.transfer_write %[[ADD]], %[[ARG0]][%[[C0]], %[[C0]]]
{{.*}} : vector<4x256xf32>, memref<4x256xf32>
// CHECK: vector.transfer_write %[[CST0]], %[[ARG0]][%[[C0]], %[[C0]]]
{{.*}} : vector<4x256xf32>, memref<4x256xf32>
___
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
[llvm-branch-commits] [mlir] 26c8f90 - [mlir[[vector] Extend Transfer read/write ops to support tensor types.
Author: Thomas Raoux
Date: 2020-12-21T08:55:04-08:00
New Revision: 26c8f9081b6b1ca9358ac2ca38e8e603fb6f7d64
URL:
https://github.com/llvm/llvm-project/commit/26c8f9081b6b1ca9358ac2ca38e8e603fb6f7d64
DIFF:
https://github.com/llvm/llvm-project/commit/26c8f9081b6b1ca9358ac2ca38e8e603fb6f7d64.diff
LOG: [mlir[[vector] Extend Transfer read/write ops to support tensor types.
Transfer_ops can now work on both buffers and tensor. Right now, lowering of
the tensor case is not supported yet.
Differential Revision: https://reviews.llvm.org/D93500
Added:
Modified:
mlir/include/mlir/Dialect/Vector/VectorOps.h
mlir/include/mlir/Dialect/Vector/VectorOps.td
mlir/include/mlir/Dialect/Vector/VectorUtils.h
mlir/include/mlir/Interfaces/VectorInterfaces.td
mlir/lib/Conversion/StandardToSPIRV/LegalizeStandardForSPIRV.cpp
mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
mlir/lib/Conversion/VectorToROCDL/VectorToROCDL.cpp
mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/lib/Dialect/Vector/VectorOps.cpp
mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp
mlir/lib/Dialect/Vector/VectorTransforms.cpp
mlir/lib/Dialect/Vector/VectorUtils.cpp
mlir/test/Dialect/Vector/invalid.mlir
mlir/test/Dialect/Vector/ops.mlir
Removed:
diff --git a/mlir/include/mlir/Dialect/Vector/VectorOps.h
b/mlir/include/mlir/Dialect/Vector/VectorOps.h
index 95964665ced6..5540a56a4043 100644
--- a/mlir/include/mlir/Dialect/Vector/VectorOps.h
+++ b/mlir/include/mlir/Dialect/Vector/VectorOps.h
@@ -126,7 +126,7 @@ namespace impl {
/// Build the default minor identity map suitable for a vector transfer. This
/// also handles the case memref<... x vector<...>> -> vector<...> in which the
/// rank of the identity map must take the vector element type into account.
-AffineMap getTransferMinorIdentityMap(MemRefType memRefType,
+AffineMap getTransferMinorIdentityMap(ShapedType shapedType,
VectorType vectorType);
} // namespace impl
} // end namespace vector
diff --git a/mlir/include/mlir/Dialect/Vector/VectorOps.td
b/mlir/include/mlir/Dialect/Vector/VectorOps.td
index de77e3b03483..13aba2076ee9 100644
--- a/mlir/include/mlir/Dialect/Vector/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/VectorOps.td
@@ -1056,7 +1056,7 @@ def Vector_TransferReadOp :
DeclareOpInterfaceMethods,
DeclareOpInterfaceMethods
]>,
-Arguments<(ins AnyMemRef:$memref, Variadic:$indices,
+Arguments<(ins AnyShaped:$source, Variadic:$indices,
AffineMapAttr:$permutation_map, AnyType:$padding,
OptionalAttr:$masked)>,
Results<(outs AnyVector:$vector)> {
@@ -1065,15 +1065,16 @@ def Vector_TransferReadOp :
let description = [{
The `vector.transfer_read` op performs a read from a slice within a
-[MemRef](../LangRef.md#memref-type) supplied as its first operand
-into a [vector](../LangRef.md#vector-type) of the same base elemental type.
+[MemRef](../LangRef.md#memref-type) or a Ranked
+[Tensor](../LangRef.md#tensor-type) supplied as its first operand into a
+[vector](../LangRef.md#vector-type) of the same base elemental type.
-A memref operand with vector element type, must have its vector element
-type match a suffix (shape and element type) of the vector (e.g.
+A memref/tensor operand with vector element type, must have its vector
+element type match a suffix (shape and element type) of the vector (e.g.
memref<3x2x6x4x3xf32>, vector<1x1x4x3xf32>).
-The slice is further defined by a full-rank index within the MemRef,
-supplied as the operands `2 .. 1 + rank(memref)`.
+The slice is further defined by a full-rank index within the MemRef/Tensor,
+supplied as the operands `2 .. 1 + rank(memref/tensor)`.
The permutation_map [attribute](../LangRef.md#attributes) is an
[affine-map](Affine.md#affine-maps) which specifies the transposition on
the
@@ -1084,8 +1085,9 @@ def Vector_TransferReadOp :
The size of the slice is specified by the size of the vector, given as the
return type.
-An `ssa-value` of the same elemental type as the MemRef is provided as the
-last operand to specify padding in the case of out-of-bounds accesses.
+An `ssa-value` of the same elemental type as the MemRef/Tensor is provided
+as the last operand to specify padding in the case of out-of-bounds
+accesses.
An optional boolean array attribute is provided to specify which dimensions
of the transfer need masking. When a dimension is specified as not
requiring
@@ -1196,17 +1198,22 @@ def Vector_TransferReadOp :
%4 = vector.transfer_read %arg1[%c3, %c3], %vf0
{permutation_map = (d0, d1)->(d0, d1)}
: memref>, v
[llvm-branch-commits] [mlir] 7c7b55b - [mlir][vector] Extend vector unroll to all element-wise ops
Author: Thomas Raoux
Date: 2020-12-21T13:31:22-08:00
New Revision: 7c7b55b985136a975223a9cefccd8fa1a5df7765
URL:
https://github.com/llvm/llvm-project/commit/7c7b55b985136a975223a9cefccd8fa1a5df7765
DIFF:
https://github.com/llvm/llvm-project/commit/7c7b55b985136a975223a9cefccd8fa1a5df7765.diff
LOG: [mlir][vector] Extend vector unroll to all element-wise ops
Extend unroll to support all element-wise ops and allow unrolling for ops with
vector operands of with the same shape as the destination but different element
type (like Cmp or Select).
Differential Revision: https://reviews.llvm.org/D93121
Added:
Modified:
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
mlir/lib/Dialect/Vector/VectorTransforms.cpp
mlir/test/Dialect/Vector/vector-transforms.mlir
mlir/test/lib/Transforms/TestVectorTransforms.cpp
Removed:
diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
index 7af44f8435ff..ba78db68214f 100644
--- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
+++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
@@ -69,7 +69,9 @@ class CastOp traits = []> :
// Base class for arithmetic cast operations.
class ArithmeticCastOp traits = []> :
-CastOp {
+CastOp])>
{
}
// Base class for unary ops. Requires single operand and result. Individual
@@ -104,6 +106,7 @@ class ArithmeticOp traits =
[]> :
Op,
ElementwiseMappable])> {
let results = (outs AnyType:$result);
@@ -992,6 +995,7 @@ def CmpFPredicateAttr : I64EnumAttr<
def CmpFOp : Std_Op<"cmpf",
[NoSideEffect, SameTypeOperands, ElementwiseMappable,
+ DeclareOpInterfaceMethods,
TypesMatchWith<
"result type has i1 element type and same shape as operands",
"lhs", "result", "getI1SameShape($_self)">]> {
@@ -1076,6 +1080,7 @@ def CmpIPredicateAttr : I64EnumAttr<
def CmpIOp : Std_Op<"cmpi",
[NoSideEffect, SameTypeOperands, ElementwiseMappable,
+ DeclareOpInterfaceMethods,
TypesMatchWith<
"result type has i1 element type and same shape as operands",
"lhs", "result", "getI1SameShape($_self)">]> {
@@ -2548,7 +2553,7 @@ def RsqrtOp : FloatUnaryOp<"rsqrt"> {
def SelectOp : Std_Op<"select", [NoSideEffect,
AllTypesMatch<["true_value", "false_value", "result"]>,
- ElementwiseMappable]> {
+ ElementwiseMappable, DeclareOpInterfaceMethods]>
{
let summary = "select operation";
let description = [{
The `select` operation chooses one value based on a binary condition
@@ -2779,7 +2784,8 @@ def SignedShiftRightOp :
IntArithmeticOp<"shift_right_signed"> {
//===--===//
def SignExtendIOp : Std_Op<"sexti",
-[NoSideEffect, ElementwiseMappable]> {
+[NoSideEffect, ElementwiseMappable,
+DeclareOpInterfaceMethods]> {
let summary = "integer sign extension operation";
let description = [{
The integer sign extension operation takes an integer input of
@@ -3595,7 +3601,9 @@ def TransposeOp : Std_Op<"transpose", [NoSideEffect]>,
// TruncateIOp
//===--===//
-def TruncateIOp : Std_Op<"trunci", [NoSideEffect, ElementwiseMappable]> {
+def TruncateIOp : Std_Op<"trunci",
+ [NoSideEffect, ElementwiseMappable,
+ DeclareOpInterfaceMethods,]> {
let summary = "integer truncation operation";
let description = [{
The integer truncation operation takes an integer input of
@@ -3862,7 +3870,9 @@ def XOrOp : IntArithmeticOp<"xor", [Commutative]> {
// ZeroExtendIOp
//===--===//
-def ZeroExtendIOp : Std_Op<"zexti", [NoSideEffect, ElementwiseMappable]> {
+def ZeroExtendIOp : Std_Op<"zexti",
+ [NoSideEffect, ElementwiseMappable,
+ DeclareOpInterfaceMethods,]> {
let summary = "integer zero extension operation";
let description = [{
The integer zero extension operation takes an integer input of
diff --git a/mlir/lib/Dialect/Vector/VectorTransforms.cpp
b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
index 1e58a759d305..5ba82b39a5a6 100644
--- a/mlir/lib/Dialect/Vector/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
@@ -492,8 +492,9 @@ static void getVectorElementwiseOpUnrollState(Operation *op,
assert(resultType && "Expected op with vector result type");
auto resultShape = resultType.getShape();
// Verify that all operands have the same vector type as result.
- assert(llvm::all_of(op->getOperandTypes(),
- [=](Type type) { return type == resultType; }));
+ assert(llvm::all_of(op->getOperandTypes(), [=](Type type) {
+return type.cast().getShape() == resultShape;
+ }));
// Create trivial elementwise identity index map based on 'resultShape'.
DenseMap indexMap;
@@ -504
[llvm-branch-commits] [mlir] 7418688 - [mlir][vector] Add more vector Ops canonicalization
Author: Thomas Raoux
Date: 2020-12-23T11:25:01-08:00
New Revision: 74186880ba99b37c0375e9d87df818beee8b4ff2
URL:
https://github.com/llvm/llvm-project/commit/74186880ba99b37c0375e9d87df818beee8b4ff2
DIFF:
https://github.com/llvm/llvm-project/commit/74186880ba99b37c0375e9d87df818beee8b4ff2.diff
LOG: [mlir][vector] Add more vector Ops canonicalization
Add canonicalization for BroadcastOp, ExtractStrideSlicesOp and ShapeCastOp
Differential Revision: https://reviews.llvm.org/D93120
Added:
Modified:
mlir/include/mlir/Dialect/Vector/VectorOps.td
mlir/lib/Dialect/Vector/VectorOps.cpp
mlir/test/Dialect/Vector/canonicalize.mlir
Removed:
diff --git a/mlir/include/mlir/Dialect/Vector/VectorOps.td
b/mlir/include/mlir/Dialect/Vector/VectorOps.td
index 13aba2076ee9..e031f87cfb8e 100644
--- a/mlir/include/mlir/Dialect/Vector/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/VectorOps.td
@@ -271,6 +271,7 @@ def Vector_BroadcastOp :
}];
let assemblyFormat = "$source attr-dict `:` type($source) `to`
type($vector)";
let hasFolder = 1;
+ let hasCanonicalizer = 1;
}
def Vector_ShuffleOp :
diff --git a/mlir/lib/Dialect/Vector/VectorOps.cpp
b/mlir/lib/Dialect/Vector/VectorOps.cpp
index a3ad355d30b2..539e00d58dbf 100644
--- a/mlir/lib/Dialect/Vector/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/VectorOps.cpp
@@ -1110,6 +1110,36 @@ OpFoldResult BroadcastOp::fold(ArrayRef
operands) {
return {};
}
+namespace {
+
+// BroadcastOp can only add dimensions or broadcast a dimension from 1 to N. In
+// the degenerated case where the broadcast only adds dimensions of size 1 it
+// can be replaced by a ShapeCastOp. This canonicalization checks if the total
+// number of elements is the same before and after the broadcast to detect if
+// the only change in the vector type are new dimensions of size 1.
+class BroadcastToShapeCast final : public OpRewritePattern {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(BroadcastOp broadcastOp,
+PatternRewriter &rewriter) const override {
+auto srcVecType = broadcastOp.getSourceType().dyn_cast();
+if (!srcVecType || broadcastOp.getVectorType().getNumElements() !=
+ srcVecType.getNumElements())
+ return failure();
+rewriter.replaceOpWithNewOp(
+broadcastOp, broadcastOp.getVectorType(), broadcastOp.source());
+return success();
+ }
+};
+
+} // namespace
+
+void BroadcastOp::getCanonicalizationPatterns(
+OwningRewritePatternList &results, MLIRContext *context) {
+ results.insert(context);
+}
+
//===--===//
// ShuffleOp
//===--===//
@@ -1768,7 +1798,8 @@ void
ExtractStridedSliceOp::getOffsets(SmallVectorImpl &results) {
namespace {
-// Pattern to rewrite a ExtractStridedSliceOp(ConstantMaskOp) ->
ConstantMaskOp.
+// Pattern to rewrite an ExtractStridedSliceOp(ConstantMaskOp) to
+// ConstantMaskOp.
class StridedSliceConstantMaskFolder final
: public OpRewritePattern {
public:
@@ -1847,14 +1878,70 @@ class StridedSliceConstantFolder final
}
};
+// Helper that returns a subset of `arrayAttr` as a vector of int64_t.
+static SmallVector getI64SubArray(ArrayAttr arrayAttr,
+ unsigned dropFront = 0,
+ unsigned dropBack = 0) {
+ assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds");
+ auto range = arrayAttr.getAsRange();
+ SmallVector res;
+ res.reserve(arrayAttr.size() - dropFront - dropBack);
+ for (auto it = range.begin() + dropFront, eit = range.end() - dropBack;
+ it != eit; ++it)
+res.push_back((*it).getValue().getSExtValue());
+ return res;
+}
+
+// Pattern to rewrite an ExtractStridedSliceOp(BroadcastOp) to
+// BroadcastOp(ExtractStrideSliceOp).
+class StridedSliceBroadcast final
+: public OpRewritePattern {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ExtractStridedSliceOp op,
+PatternRewriter &rewriter) const override {
+auto broadcast = op.vector().getDefiningOp();
+if (!broadcast)
+ return failure();
+auto srcVecType = broadcast.source().getType().dyn_cast();
+unsigned srcRrank = srcVecType ? srcVecType.getRank() : 0;
+auto dstVecType = op.getType().cast();
+unsigned dstRank = dstVecType.getRank();
+unsigned rankDiff = dstRank - srcRrank;
+// Check if the most inner dimensions of the source of the broacast are the
+// same as the destination of the extract. If this is the case we can just
+// use a broadcast as the original dimensions are untouched.
+bool lowerDimMatch = true;
+for (unsigned i = 0; i < srcRrank; i++) {
+
[llvm-branch-commits] [mlir] c503dc1 - [mlir][linalg] Add vectorization for element-wise linalg ops
Author: Thomas Raoux
Date: 2020-12-03T15:31:13-08:00
New Revision: c503dc1b8a52946e4daefa1a266e74a102382971
URL:
https://github.com/llvm/llvm-project/commit/c503dc1b8a52946e4daefa1a266e74a102382971
DIFF:
https://github.com/llvm/llvm-project/commit/c503dc1b8a52946e4daefa1a266e74a102382971.diff
LOG: [mlir][linalg] Add vectorization for element-wise linalg ops
Add support for vectorization for linalg.generic representing element-wise ops.
Those are converted to transfer_read + vector ops + transfer_write.
Also re-organize the vectorization tests to be together.
Implementation derived from the work of @burmako, @agrue and
@fedelebron.
Differential Revision: https://reviews.llvm.org/D92540
Added:
mlir/test/Dialect/Linalg/vectorization.mlir
Modified:
mlir/include/mlir/EDSC/Builders.h
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/lib/EDSC/Builders.cpp
mlir/test/Dialect/Linalg/transform-patterns-matmul-to-vector.mlir
mlir/test/Dialect/Linalg/transform-patterns.mlir
mlir/test/lib/Transforms/TestLinalgTransforms.cpp
Removed:
diff --git a/mlir/include/mlir/EDSC/Builders.h
b/mlir/include/mlir/EDSC/Builders.h
index 70c948d99cda..83b6634bf8e2 100644
--- a/mlir/include/mlir/EDSC/Builders.h
+++ b/mlir/include/mlir/EDSC/Builders.h
@@ -30,6 +30,7 @@ namespace edsc {
/// setting and restoring of insertion points.
class ScopedContext {
public:
+ ScopedContext(OpBuilder &b);
ScopedContext(OpBuilder &b, Location location);
/// Sets the insertion point of the builder to 'newInsertPt' for the duration
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 8860674ef847..a28b90b1d95c 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -84,6 +84,195 @@ static LogicalResult isContraction(Operation *op) {
hasMultiplyAddBody(genericOp.region()));
}
+static bool hasOnlyScalarElementwiseOp(Region &r) {
+ if (!llvm::hasSingleElement(r))
+return false;
+ for (Operation &op : r.front()) {
+if (!(isa(op) ||
+ op.hasTrait()) ||
+llvm::any_of(op.getResultTypes(),
+ [](Type type) { return !type.isIntOrIndexOrFloat(); }))
+ return false;
+ }
+ return true;
+}
+
+// Return true if the op is an element-wise linalg op.
+static bool isElementwise(Operation *op) {
+ auto genericOp = dyn_cast(op);
+ if (!genericOp)
+return false;
+ if (genericOp.getNumLoops() != genericOp.getNumParallelLoops())
+return false;
+ // TODO: relax the restrictions on indexing map.
+ for (unsigned i = 0, e = genericOp.getNumOutputs(); i < e; i++) {
+if (!genericOp.getOutputIndexingMap(i).isIdentity())
+ return false;
+ }
+ // Currently limit the input indexing map to minor identity as other
+ // permutations might require adding transpose ops to convert the vector read
+ // to the right shape.
+ for (unsigned i = 0, e = genericOp.getNumInputs(); i < e; i++) {
+if (!genericOp.getInputIndexingMap(i).isMinorIdentity())
+ return false;
+ }
+ return hasOnlyScalarElementwiseOp(genericOp.getRegion());
+}
+
+static VectorType extractVectorTypeFromScalarView(Value v) {
+ MemRefType mt = v.getType().cast();
+ return mt.getShape().empty()
+ ? VectorType()
+ : VectorType::get(mt.getShape(), mt.getElementType());
+}
+
+static Value transferReadVector(OpBuilder &builder, Value memref) {
+ edsc::ScopedContext scope(builder);
+ auto memrefType = memref.getType().cast();
+ if (VectorType vectorType = extractVectorTypeFromScalarView(memref)) {
+SmallVector indices(memrefType.getRank(), std_constant_index(0));
+return vector_transfer_read(vectorType, memref, indices);
+ }
+ return std_load(memref);
+}
+
+static void transferWriteVector(OpBuilder &builder, Value value, Value memref)
{
+ edsc::ScopedContext scope(builder);
+ auto memrefType = memref.getType().cast();
+ if (VectorType vectorType = extractVectorTypeFromScalarView(memref)) {
+SmallVector indices(memrefType.getRank(), std_constant_index(0));
+if (vectorType != value.getType())
+ value = vector_broadcast(vectorType, value);
+vector_transfer_write(value, memref, indices);
+ } else {
+std_store(value, memref);
+ }
+}
+
+namespace {
+// Transforms scalar operations into their vectorized counterparts,
+// while using the provided generic op to map:
+// * Its arguments to transfer reads from the views of the generic op.
+// * linalg.yield ops to transfer writes to the views of the generic op.
+class GenericVectorizer {
+public:
+ GenericVectorizer(OpBuilder &builder, linalg::GenericOp generic)
+ : builder(builder), generic(generic) {}
+
+ // Takes a scalar operation and builds its vectorized counterpart or
+ // counterparts using the underlying builder.
+ //
[llvm-branch-commits] [mlir] 3e3e276 - [mlir][vector][NFC] Change UnrollVectorPattern to not be statically dependent on an op type
Author: Thomas Raoux
Date: 2020-12-04T09:53:01-08:00
New Revision: 3e3e276d22ca6917a721c4173b00b37850d8020c
URL:
https://github.com/llvm/llvm-project/commit/3e3e276d22ca6917a721c4173b00b37850d8020c
DIFF:
https://github.com/llvm/llvm-project/commit/3e3e276d22ca6917a721c4173b00b37850d8020c.diff
LOG: [mlir][vector][NFC] Change UnrollVectorPattern to not be statically
dependent on an op type
Make UnrollVectorPattern inherit from RewritePattern instead of
OpRewritePattern so that we don't need to create many patterns when applying to
many different type of ops. Since we may want to apply the pattern to all
arithmetic op, it is more convenient to filter dynamically.
Differential Revision: https://reviews.llvm.org/D92635
Added:
Modified:
mlir/include/mlir/Dialect/Vector/VectorTransforms.h
mlir/test/lib/Transforms/TestVectorTransforms.cpp
Removed:
diff --git a/mlir/include/mlir/Dialect/Vector/VectorTransforms.h
b/mlir/include/mlir/Dialect/Vector/VectorTransforms.h
index cc0b2841d3f1..c88aa7f5bc65 100644
--- a/mlir/include/mlir/Dialect/Vector/VectorTransforms.h
+++ b/mlir/include/mlir/Dialect/Vector/VectorTransforms.h
@@ -91,7 +91,7 @@ struct UnrollVectorOptions {
/// Callback function that indicates whether vector unrolling should be
/// attempted on the operation.
FilterConstraintFnType filterConstraint = nullptr;
- UnrollVectorOptions &setFilterContraint(FilterConstraintFnType constraint) {
+ UnrollVectorOptions &setFilterConstraint(FilterConstraintFnType constraint) {
filterConstraint = constraint;
return *this;
}
@@ -117,21 +117,19 @@ struct UnrollVectorOptions {
};
/// Pattern to apply `unrollSingleResultVectorOp` to a `targetShape`
/// declaratively.
-template
-struct UnrollVectorPattern : public OpRewritePattern {
- using FilterConstraintType = std::function;
+struct UnrollVectorPattern : public RewritePattern {
+ using FilterConstraintType = std::function;
UnrollVectorPattern(MLIRContext *context, UnrollVectorOptions options)
- : OpRewritePattern(context), options(options) {}
- LogicalResult matchAndRewrite(OpTy op,
+ : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()), options(options) {}
+ LogicalResult matchAndRewrite(Operation *op,
PatternRewriter &rewriter) const override {
if (options.filterConstraint && failed(options.filterConstraint(op)))
return failure();
if (!options.nativeShape) {
- return op.emitError("vector unrolling expects the native shape or native"
- "shape call back function to be set");
+ return op->emitError("vector unrolling expects the native shape or
native"
+ "shape call back function to be set");
}
-auto unrollableVectorOp =
-dyn_cast(op.getOperation());
+auto unrollableVectorOp = dyn_cast(op);
if (!unrollableVectorOp)
return failure();
auto maybeUnrollShape = unrollableVectorOp.getShapeForUnroll();
@@ -139,12 +137,12 @@ struct UnrollVectorPattern : public
OpRewritePattern {
return failure();
Optional> targetShape = options.nativeShape(op);
if (!targetShape)
- return op.emitError("failed to get target shape for vector unroll");
+ return op->emitError("failed to get target shape for vector unroll");
auto maybeShapeRatio = shapeRatio(*maybeUnrollShape, *targetShape);
if (!maybeShapeRatio ||
llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; }))
return failure();
-if (std::is_same::value) {
+if (isa(op)) {
if (failed(unrollTransferWriteOp(rewriter, op, *targetShape)))
return failure();
rewriter.eraseOp(op);
diff --git a/mlir/test/lib/Transforms/TestVectorTransforms.cpp
b/mlir/test/lib/Transforms/TestVectorTransforms.cpp
index 602bf8148cd8..99c336ef0565 100644
--- a/mlir/test/lib/Transforms/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Transforms/TestVectorTransforms.cpp
@@ -27,14 +27,22 @@ struct TestVectorToVectorConversion
void runOnFunction() override {
OwningRewritePatternList patterns;
auto *ctx = &getContext();
-patterns.insert>(
-ctx, UnrollVectorOptions().setNativeShape(ArrayRef{2, 2}));
-patterns.insert>(
-ctx, UnrollVectorOptions().setNativeShape(ArrayRef{2, 2, 2}));
+patterns.insert(
+ctx, UnrollVectorOptions().setNativeShapeFn(getShape));
populateVectorToVectorCanonicalizationPatterns(patterns, ctx);
populateVectorToVectorTransformationPatterns(patterns, ctx);
applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
}
+
+private:
+ // Return the target shape based on op type.
+ static Optional> getShape(Operation *op) {
+if (isa(op))
+ return SmallVector(2, 2);
+if (isa(op))
+ return SmallVector(3, 2);
+return llvm::None;
+ }
};
struct TestVectorSlicesConversion
@@
