[MLIR] Add support for permutation_map
This CL hooks up and uses permutation_map in vector_transfer ops.
In particular, when going into the nuts and bolts of the implementation, it
became clear that cases arose that required supporting broadcast semantics.
Broadcast semantics are thus added to the general permutation_map.
The verify methods and tests are updated accordingly.
Examples of interest include.
Example 1:
The following MLIR snippet:
```mlir
for %i3 = 0 to %M {
for %i4 = 0 to %N {
for %i5 = 0 to %P {
%a5 = load %A[%i4, %i5, %i3] : memref<?x?x?xf32>
}}}
```
may vectorize with {permutation_map: (d0, d1, d2) -> (d2, d1)} into:
```mlir
for %i3 = 0 to %0 step 32 {
for %i4 = 0 to %1 {
for %i5 = 0 to %2 step 256 {
%4 = vector_transfer_read %arg0, %i4, %i5, %i3
{permutation_map: (d0, d1, d2) -> (d2, d1)} :
(memref<?x?x?xf32>, index, index) -> vector<32x256xf32>
}}}
````
Meaning that vector_transfer_read will be responsible for reading the 2-D slice:
`%arg0[%i4, %i5:%15+256, %i3:%i3+32]` into vector<32x256xf32>. This will
require a transposition when vector_transfer_read is further lowered.
Example 2:
The following MLIR snippet:
```mlir
%cst0 = constant 0 : index
for %i0 = 0 to %M {
%a0 = load %A[%cst0, %cst0] : memref<?x?xf32>
}
```
may vectorize with {permutation_map: (d0) -> (0)} into:
```mlir
for %i0 = 0 to %0 step 128 {
%3 = vector_transfer_read %arg0, %c0_0, %c0_0
{permutation_map: (d0, d1) -> (0)} :
(memref<?x?xf32>, index, index) -> vector<128xf32>
}
````
Meaning that vector_transfer_read will be responsible of reading the 0-D slice
`%arg0[%c0, %c0]` into vector<128xf32>. This will require a 1-D vector
broadcast when vector_transfer_read is further lowered.
Additionally, some minor cleanups and refactorings are performed.
One notable thing missing here is the composition with a projection map during
materialization. This is because I could not find an AffineMap composition
that operates on AffineMap directly: everything related to composition seems
to require going through SSAValue and only operates on AffinMap at a distance
via AffineValueMap. I have raised this concern a bunch of times already, the
followup CL will actually do something about it.
In the meantime, the projection is hacked at a minimum to pass verification
and materialiation tests are temporarily incorrect.
PiperOrigin-RevId: 224376828
diff --git a/include/mlir/Analysis/LoopAnalysis.h b/include/mlir/Analysis/LoopAnalysis.h
index 5900a4c..db401a4 100644
--- a/include/mlir/Analysis/LoopAnalysis.h
+++ b/include/mlir/Analysis/LoopAnalysis.h
@@ -22,12 +22,15 @@
#ifndef MLIR_ANALYSIS_LOOP_ANALYSIS_H
#define MLIR_ANALYSIS_LOOP_ANALYSIS_H
+#include "mlir/Support/LLVM.h"
+
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/Optional.h"
namespace mlir {
class AffineExpr;
+class AffineMap;
class ForStmt;
class MemRefType;
class MLValue;
@@ -48,12 +51,29 @@
/// this method is thus able to determine non-trivial divisors.
uint64_t getLargestDivisorOfTripCount(const ForStmt &forStmt);
-/// Given a MemRef accessed by `indices` and a dimension `dim`, determines
-/// whether indices[dim] is independent of the value `input`.
-// For now we assume no layout map or identity layout map in the MemRef.
-// TODO(ntv): support more than identity layout map.
-bool isAccessInvariant(const MLValue &input, MemRefType memRefType,
- llvm::ArrayRef<const MLValue *> indices, unsigned dim);
+/// Given an induction variable `iv` of type ForStmt and an `index` of type
+/// IndexType, returns `true` if `index` is independent of `iv` and false
+/// otherwise.
+/// The determination supports composition with at most one AffineApplyOp.
+/// The at most one AffineApplyOp comes from the fact that composition of
+/// AffineApplyOp need to be canonicalized by construction to avoid writing code
+/// that composes arbitrary numbers of AffineApplyOps everywhere. To achieve
+/// this, at the very least, the compose-affine-apply pass must have been run.
+///
+/// Prerequisites:
+/// 1. `iv` and `index` of the proper type;
+/// 2. at most one reachable AffineApplyOp from index;
+bool isAccessInvariant(const MLValue &iv, const MLValue &index);
+
+/// Given an induction variable `iv` of type ForStmt and `indices` of type
+/// IndexType, returns the set of `indices` that are independent of `iv`.
+///
+/// Prerequisites (inherited from `isAccessInvariant` above):
+/// 1. `iv` and `indices` of the proper type;
+/// 2. at most one reachable AffineApplyOp from index;
+llvm::DenseSet<const MLValue *, llvm::DenseMapInfo<const MLValue *>>
+getInvariantAccesses(const MLValue &iv,
+ llvm::ArrayRef<const MLValue *> indices);
/// Checks whether the loop is structurally vectorizable; i.e.:
/// 1. the loop has proper dependence semantics (parallel, reduction, etc);
diff --git a/include/mlir/Analysis/VectorAnalysis.h b/include/mlir/Analysis/VectorAnalysis.h
index a3d31b2..e37043f 100644
--- a/include/mlir/Analysis/VectorAnalysis.h
+++ b/include/mlir/Analysis/VectorAnalysis.h
@@ -20,9 +20,12 @@
#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/DenseMap.h"
+
namespace mlir {
class AffineMap;
+class ForStmt;
class MemRefType;
class OperationStmt;
class VectorType;
@@ -45,15 +48,74 @@
llvm::Optional<llvm::SmallVector<unsigned, 4>>
shapeRatio(VectorType superVectorType, VectorType subVectorType);
-/// Creates a permutation map to be used as an attribute in VectorTransfer ops.
-/// Currently only returns the minor vectorType.rank identity submatrix.
+/// Constructs a permutation map of invariant memref indices to vector
+/// dimension.
///
-/// For example, assume memrefType is of rank 5 and vectorType is of rank 3,
-/// returns the affine map:
-/// (d0, d1, d2, d3, d4) -> (d2, d3, d4)
+/// If no index is found to be invariant, 0 is added to the permutation_map and
+/// corresponds to a vector broadcast along that dimension.
///
-/// TODO(ntv): support real permutations.
-AffineMap makePermutationMap(MemRefType memrefType, VectorType vectorType);
+/// The implementation uses the knowledge of the mapping of loops to
+/// vector dimension. `loopToVectorDim` carries this information as a map with:
+/// - keys representing "vectorized enclosing loops";
+/// - values representing the corresponding vector dimension.
+/// Note that loopToVectorDim is a whole function map from which only enclosing
+/// loop information is extracted.
+///
+/// Prerequisites: `opStmt` is a vectorizable load or store operation (i.e. at
+/// most one invariant index along each ForStmt of `loopToVectorDim`).
+///
+/// Example 1:
+/// The following MLIR snippet:
+///
+/// ```mlir
+/// for %i3 = 0 to %0 {
+/// for %i4 = 0 to %1 {
+/// for %i5 = 0 to %2 {
+/// %a5 = load %arg0[%i4, %i5, %i3] : memref<?x?x?xf32>
+/// }}}
+/// ```
+///
+/// may vectorize with {permutation_map: (d0, d1, d2) -> (d2, d1)} into:
+///
+/// ```mlir
+/// for %i3 = 0 to %0 step 32 {
+/// for %i4 = 0 to %1 {
+/// for %i5 = 0 to %2 step 256 {
+/// %4 = vector_transfer_read %arg0, %i4, %i5, %i3
+/// {permutation_map: (d0, d1, d2) -> (d2, d1)} :
+/// (memref<?x?x?xf32>, index, index) -> vector<32x256xf32>
+/// }}}
+/// ```
+///
+/// Meaning that vector_transfer_read will be responsible for reading the slice:
+/// `%arg0[%i4, %i5:%15+256, %i3:%i3+32]` into vector<32x256xf32>.
+///
+/// Example 2:
+/// The following MLIR snippet:
+///
+/// ```mlir
+/// %cst0 = constant 0 : index
+/// for %i0 = 0 to %0 {
+/// %a0 = load %arg0[%cst0, %cst0] : memref<?x?xf32>
+/// }
+/// ```
+///
+/// may vectorize with {permutation_map: (d0) -> (0)} into:
+///
+/// ```mlir
+/// for %i0 = 0 to %0 step 128 {
+/// %3 = vector_transfer_read %arg0, %c0_0, %c0_0
+/// {permutation_map: (d0, d1) -> (0)} :
+/// (memref<?x?xf32>, index, index) -> vector<128xf32>
+/// }
+/// ````
+///
+/// Meaning that vector_transfer_read will be responsible of reading the slice
+/// `%arg0[%c0, %c0]` into vector<128xf32> which needs a 1-D vector broadcast.
+///
+AffineMap
+makePermutationMap(OperationStmt *opStmt,
+ const llvm::DenseMap<ForStmt *, unsigned> &loopToVectorDim);
namespace matcher {
diff --git a/include/mlir/StandardOps/StandardOps.h b/include/mlir/StandardOps/StandardOps.h
index 562c932..17c3da8 100644
--- a/include/mlir/StandardOps/StandardOps.h
+++ b/include/mlir/StandardOps/StandardOps.h
@@ -859,12 +859,19 @@
/// vector read exceeds the MemRef bounds. If the value is not specified,
/// the access is statically guaranteed to be within bounds;
/// 2. an attribute of type AffineMap to specify a slice of the original
-/// MemRef access and its transposition into the super-vector shape. The
-/// permutation_map is an unbounded AffineMap that must represent a
-/// permutation from the MemRef dim space projected onto the vector dim
-/// space.
-//
-/// Example:
+/// MemRef access and its transposition into the super-vector shape.
+/// The permutation_map is an unbounded AffineMap that must
+/// represent a permutation from the MemRef dim space projected onto the
+/// vector dim space.
+/// This permutation_map has as many output dimensions as the vector rank.
+/// However, it is not necessarily full rank on the target space to signify
+/// that broadcast operations will be needed along certain vector
+/// dimensions.
+/// In the limit, one may load a 0-D slice of a memref (i.e. a single
+/// value) into a vector, which corresponds to broadcasting that value in
+/// the whole vector (i.e. a non-constant splat).
+///
+/// Example with full rank permutation_map:
/// ```mlir
/// %A = alloc(%size1, %size2, %size3, %size4) : memref<?x?x?x?xf32>
/// ...
@@ -879,6 +886,17 @@
/// (memref<?x?x?x?xf32>, index, index, index, index, f32) ->
/// vector<16x32x64xf32>
/// ```
+///
+/// Example with partial rank permutation_map:
+/// ```mlir
+/// %c0 = constant 0 : index
+/// %A = alloc(%size1, %size2, %size3, %size4) : memref<?x?x?x?xf32>
+/// ...
+/// // let %i, %j be ssa-values of type index
+/// %v0 = vector_transfer_read %src, %i, %c0, %c0, %c0
+/// {permutation_map: (d0, d1, d2, d3) -> (0, d1, 0)} :
+/// (memref<?x?x?x?xf32>, index, index, index, index) ->
+/// vector<16x32x64xf32>
class VectorTransferReadOp
: public Op<VectorTransferReadOp, OpTrait::VariadicOperands,
OpTrait::OneResult> {
@@ -898,6 +916,7 @@
const SSAValue *getMemRef() const {
return getOperand(Offsets::MemRefOffset);
}
+ VectorType getVectorType() const { return getResultType(); }
MemRefType getMemRefType() const {
return getMemRef()->getType().cast<MemRefType>();
}
diff --git a/include/mlir/Support/LLVM.h b/include/mlir/Support/LLVM.h
index d3f3cd5..2f60d59 100644
--- a/include/mlir/Support/LLVM.h
+++ b/include/mlir/Support/LLVM.h
@@ -52,6 +52,7 @@
template <typename KeyT, typename ValueT> struct DenseMapPair;
}
template <typename T> struct DenseMapInfo;
+template <typename ValueT, typename ValueInfoT> class DenseSet;
template <typename KeyT, typename ValueT, typename KeyInfoT, typename BucketT>
class DenseMap;
@@ -72,6 +73,8 @@
// Containers.
using llvm::ArrayRef;
using llvm::DenseMap;
+using llvm::DenseMapInfo;
+using llvm::DenseSet;
using llvm::MutableArrayRef;
using llvm::None;
using llvm::Optional;
diff --git a/lib/Analysis/LoopAnalysis.cpp b/lib/Analysis/LoopAnalysis.cpp
index de98849..253d544 100644
--- a/lib/Analysis/LoopAnalysis.cpp
+++ b/lib/Analysis/LoopAnalysis.cpp
@@ -31,7 +31,10 @@
#include "mlir/StandardOps/StandardOps.h"
#include "mlir/Support/Functional.h"
#include "mlir/Support/MathExtras.h"
+
+#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SmallString.h"
+#include <type_traits>
using namespace mlir;
@@ -120,65 +123,91 @@
return tripCountExpr.getLargestKnownDivisor();
}
-bool mlir::isAccessInvariant(const MLValue &input, MemRefType memRefType,
- ArrayRef<const MLValue *> indices, unsigned dim) {
- assert(indices.size() == memRefType.getRank());
- assert(dim < indices.size());
- auto layoutMap = memRefType.getAffineMaps();
- assert(memRefType.getAffineMaps().size() <= 1);
- // TODO(ntv): remove dependence on Builder once we support non-identity
- // layout map.
- Builder b(memRefType.getContext());
- assert(layoutMap.empty() ||
- layoutMap[0] == b.getMultiDimIdentityMap(indices.size()));
- (void)layoutMap;
-
+bool mlir::isAccessInvariant(const MLValue &iv, const MLValue &index) {
+ assert(isa<ForStmt>(iv) && "iv must be a ForStmt");
+ assert(index.getType().isa<IndexType>() && "index must be of IndexType");
SmallVector<OperationStmt *, 4> affineApplyOps;
- getReachableAffineApplyOps({const_cast<MLValue *>(indices[dim])},
- affineApplyOps);
+ getReachableAffineApplyOps({const_cast<MLValue *>(&index)}, affineApplyOps);
if (affineApplyOps.empty()) {
// Pointer equality test because of MLValue pointer semantics.
- return indices[dim] != &input;
+ return &index != &iv;
}
- assert(affineApplyOps.size() == 1 &&
- "CompositionAffineMapsPass must have "
- "been run: there should be at most one AffineApplyOp");
+ assert(
+ affineApplyOps.size() == 1 &&
+ "CompositionAffineMapsPass must have been run: there should be at most "
+ "one AffineApplyOp");
+
auto composeOp = affineApplyOps[0]->cast<AffineApplyOp>();
// We need yet another level of indirection because the `dim` index of the
// access may not correspond to the `dim` index of composeOp.
unsigned idx = std::numeric_limits<unsigned>::max();
unsigned numResults = composeOp->getNumResults();
for (unsigned i = 0; i < numResults; ++i) {
- if (indices[dim] == composeOp->getResult(i)) {
+ if (&index == composeOp->getResult(i)) {
idx = i;
break;
}
}
assert(idx < std::numeric_limits<unsigned>::max());
return !AffineValueMap(*composeOp)
- .isFunctionOf(idx, &const_cast<MLValue &>(input));
+ .isFunctionOf(idx, &const_cast<MLValue &>(iv));
}
-/// Determines whether a load or a store has a contiguous access along the
-/// value `input`. Contiguous is defined as either invariant or varying only
-/// along the fastest varying memory dimension.
-// TODO(ntv): allow more advanced notions of contiguity (non-fastest varying,
-// check strides, ...).
-template <typename LoadOrStoreOpPointer>
-static bool isContiguousAccess(const MLValue &input,
- LoadOrStoreOpPointer memoryOp,
+llvm::DenseSet<const MLValue *>
+mlir::getInvariantAccesses(const MLValue &iv,
+ llvm::ArrayRef<const MLValue *> indices) {
+ llvm::DenseSet<const MLValue *> res;
+ for (unsigned idx = 0, n = indices.size(); idx < n; ++idx) {
+ auto *val = indices[idx];
+ if (isAccessInvariant(iv, *val)) {
+ res.insert(val);
+ }
+ }
+ return res;
+}
+
+/// Given:
+/// 1. an induction variable `iv` of type ForStmt;
+/// 2. a `memoryOp` of type const LoadOp& or const StoreOp&;
+/// 3. the index of the `fastestVaryingDim` along which to check;
+/// determines whether `memoryOp`[`fastestVaryingDim`] is a contiguous access
+/// along `iv`.
+/// Contiguous is defined as either invariant or varying only along
+/// `fastestVaryingDim`.
+///
+/// Prerequisites:
+/// 1. `iv` of the proper type;
+/// 2. the MemRef accessed by `memoryOp` has no layout map or at most an
+/// identity layout map.
+///
+// TODO(ntv): check strides.
+template <typename LoadOrStoreOp>
+static bool isContiguousAccess(const MLValue &iv, const LoadOrStoreOp &memoryOp,
unsigned fastestVaryingDim) {
- using namespace functional;
- auto indices = map([](const SSAValue *val) { return dyn_cast<MLValue>(val); },
- memoryOp->getIndices());
- auto memRefType = memoryOp->getMemRefType();
- for (unsigned d = 0, numIndices = indices.size(); d < numIndices; ++d) {
- if (fastestVaryingDim == (numIndices - 1) - d) {
+ static_assert(std::is_same<LoadOrStoreOp, LoadOp>::value ||
+ std::is_same<LoadOrStoreOp, StoreOp>::value,
+ "Must be called on either const LoadOp & or const StoreOp &");
+ auto memRefType = memoryOp.getMemRefType();
+ auto layoutMap = memRefType.getAffineMaps();
+ (void)layoutMap;
+ Builder b(memoryOp.getOperation()->getContext());
+ (void)b;
+ assert(layoutMap.empty() ||
+ (layoutMap.size() == 1 &&
+ layoutMap[0] == b.getMultiDimIdentityMap(layoutMap[0].getNumDims())));
+ assert(fastestVaryingDim < memRefType.getRank());
+
+ auto indices = memoryOp.getIndices();
+ // TODO(clattner): should iterator_range have a size method?
+ auto numIndices = indices.end() - indices.begin();
+ unsigned d = 0;
+ for (auto index : indices) {
+ if (fastestVaryingDim == (numIndices - 1) - d++) {
continue;
}
- if (!isAccessInvariant(input, memRefType, indices, d)) {
+ if (!isAccessInvariant(iv, cast<MLValue>(*index))) {
return false;
}
}
@@ -247,8 +276,8 @@
[fastestVaryingDim](const ForStmt &loop, const OperationStmt &op) {
auto load = op.dyn_cast<LoadOp>();
auto store = op.dyn_cast<StoreOp>();
- return load ? isContiguousAccess(loop, load, fastestVaryingDim)
- : isContiguousAccess(loop, store, fastestVaryingDim);
+ return load ? isContiguousAccess(loop, *load, fastestVaryingDim)
+ : isContiguousAccess(loop, *store, fastestVaryingDim);
});
return isVectorizableLoopWithCond(loop, fun);
}
diff --git a/lib/Analysis/VectorAnalysis.cpp b/lib/Analysis/VectorAnalysis.cpp
index 9c2160c..ebddaff 100644
--- a/lib/Analysis/VectorAnalysis.cpp
+++ b/lib/Analysis/VectorAnalysis.cpp
@@ -16,12 +16,17 @@
// =============================================================================
#include "mlir/Analysis/VectorAnalysis.h"
+#include "mlir/Analysis/LoopAnalysis.h"
+#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Statements.h"
#include "mlir/StandardOps/StandardOps.h"
#include "mlir/Support/Functional.h"
#include "mlir/Support/STLExtras.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SetVector.h"
+
///
/// Implements Analysis functions specific to vectors which support
/// the vectorization and vectorization materialization passes.
@@ -29,6 +34,8 @@
using namespace mlir;
+using llvm::SetVector;
+
Optional<SmallVector<unsigned, 4>> mlir::shapeRatio(ArrayRef<int> superShape,
ArrayRef<int> subShape) {
if (superShape.size() < subShape.size()) {
@@ -76,18 +83,98 @@
return shapeRatio(superVectorType.getShape(), subVectorType.getShape());
}
-AffineMap mlir::makePermutationMap(MemRefType memrefType,
- VectorType vectorType) {
- unsigned memRefRank = memrefType.getRank();
- unsigned vectorRank = vectorType.getRank();
- assert(memRefRank >= vectorRank && "Broadcast not supported");
- unsigned offset = memRefRank - vectorRank;
- SmallVector<AffineExpr, 4> perm;
- perm.reserve(memRefRank);
- for (unsigned i = 0; i < vectorRank; ++i) {
- perm.push_back(getAffineDimExpr(offset + i, memrefType.getContext()));
+/// Constructs a permutation map from memref indices to vector dimension.
+///
+/// The implementation uses the knowledge of the mapping of enclosing loop to
+/// vector dimension. `enclosingLoopToVectorDim` carries this information as a
+/// map with:
+/// - keys representing "vectorized enclosing loops";
+/// - values representing the corresponding vector dimension.
+/// The algorithm traverses "vectorized enclosing loops" and extracts the
+/// at-most-one MemRef index that is invariant along said loop. This index is
+/// guaranteed to be at most one by construction: otherwise the MemRef is not
+/// vectorizable.
+/// If this invariant index is found, it is added to the permutation_map at the
+/// proper vector dimension.
+/// If no index is found to be invariant, 0 is added to the permutation_map and
+/// corresponds to a vector broadcast along that dimension.
+///
+/// Examples can be found in the documentation of `makePermutationMap`, in the
+/// header file.
+static AffineMap makePermutationMap(
+ MLIRContext *context,
+ llvm::iterator_range<Operation::operand_iterator> indices,
+ const DenseMap<ForStmt *, unsigned> &enclosingLoopToVectorDim) {
+ using functional::makePtrDynCaster;
+ using functional::map;
+ auto unwrappedIndices = map(makePtrDynCaster<SSAValue, MLValue>(), indices);
+ SmallVector<AffineExpr, 4> perm(enclosingLoopToVectorDim.size(),
+ getAffineConstantExpr(0, context));
+ for (auto kvp : enclosingLoopToVectorDim) {
+ assert(kvp.second < perm.size());
+ auto invariants = getInvariantAccesses(*kvp.first, unwrappedIndices);
+ unsigned numIndices = unwrappedIndices.size();
+ unsigned countInvariantIndices = 0;
+ for (unsigned dim = 0; dim < numIndices; ++dim) {
+ if (!invariants.count(unwrappedIndices[dim])) {
+ assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
+ "permutationMap already has an entry along dim");
+ perm[kvp.second] = getAffineDimExpr(dim, context);
+ } else {
+ ++countInvariantIndices;
+ }
+ }
+ assert((countInvariantIndices == numIndices ||
+ countInvariantIndices == numIndices - 1) &&
+ "Vectorization prerequisite violated: at most 1 index may be "
+ "invariant wrt a vectorized loop");
}
- return AffineMap::get(memRefRank, 0, perm, {});
+ return AffineMap::get(unwrappedIndices.size(), 0, perm, {});
+}
+
+/// Implementation detail that walks up the parents and records the ones with
+/// the specified type.
+/// TODO(ntv): could also be implemented as a collect parents followed by a
+/// filter and made available outside this file.
+template <typename T> static SetVector<T *> getParentsOfType(Statement *stmt) {
+ SetVector<T *> res;
+ auto *current = stmt;
+ while (auto *parent = current->getParentStmt()) {
+ auto *typedParent = dyn_cast<T>(parent);
+ if (typedParent) {
+ assert(res.count(typedParent) == 0 && "Already inserted");
+ res.insert(typedParent);
+ }
+ current = parent;
+ }
+ return res;
+}
+
+/// Returns the enclosing ForStmt, from closest to farthest.
+static SetVector<ForStmt *> getEnclosingForStmts(Statement *stmt) {
+ return getParentsOfType<ForStmt>(stmt);
+}
+
+AffineMap
+mlir::makePermutationMap(OperationStmt *opStmt,
+ const DenseMap<ForStmt *, unsigned> &loopToVectorDim) {
+ DenseMap<ForStmt *, unsigned> enclosingLoopToVectorDim;
+ auto enclosingLoops = getEnclosingForStmts(opStmt);
+ for (auto *forStmt : enclosingLoops) {
+ auto it = loopToVectorDim.find(forStmt);
+ if (it != loopToVectorDim.end()) {
+ enclosingLoopToVectorDim.insert(*it);
+ }
+ }
+
+ if (auto load = opStmt->dyn_cast<LoadOp>()) {
+ return ::makePermutationMap(opStmt->getContext(), load->getIndices(),
+ enclosingLoopToVectorDim);
+ }
+
+ auto store = opStmt->cast<StoreOp>();
+ return ::makePermutationMap(opStmt->getContext(), store->getIndices(),
+ enclosingLoopToVectorDim);
}
bool mlir::matcher::operatesOnStrictSuperVectors(const OperationStmt &opStmt,
diff --git a/lib/StandardOps/StandardOps.cpp b/lib/StandardOps/StandardOps.cpp
index 4f83754..e125c62 100644
--- a/lib/StandardOps/StandardOps.cpp
+++ b/lib/StandardOps/StandardOps.cpp
@@ -1374,15 +1374,23 @@
SmallVector<bool, 8> seen(permutationMap.getNumInputs(), false);
for (auto expr : permutationMap.getResults()) {
auto dim = expr.dyn_cast<AffineDimExpr>();
+ auto zero = expr.dyn_cast<AffineConstantExpr>();
+ if (zero) {
+ if (zero.getValue() != 0) {
+ return emitOpError(
+ "requires a projected permutation_map (at most one dim or the zero "
+ "constant can appear in each result)");
+ }
+ continue;
+ }
if (!dim) {
- return emitOpError(
- "requires a permutation_map that is an actual permutation");
+ return emitOpError("requires a projected permutation_map (at most one "
+ "dim or the zero constant can appear in each result)");
}
if (seen[dim.getPosition()]) {
return emitOpError(
- "requires a permutation_map that is a full column-rank "
- "permutation (i.e. a permutation composed with an "
- "orthogonal projection)");
+ "requires a permutation_map that is a permutation (found one dim "
+ "used more than once)");
}
seen[dim.getPosition()] = true;
}
diff --git a/lib/Transforms/MaterializeVectors.cpp b/lib/Transforms/MaterializeVectors.cpp
index 27f157c..0d7d0db 100644
--- a/lib/Transforms/MaterializeVectors.cpp
+++ b/lib/Transforms/MaterializeVectors.cpp
@@ -82,6 +82,73 @@
/// operations and builds the slice scoped the innermost loop enclosing the
/// current vector_transfer_write. These assumptions and the implementation
/// details are subject to revision in the future.
+///
+/// Example
+/// ========
+/// In the following, the single vector_transfer_write op operates on a
+/// vector<4x4x4xf32>. Let's assume the HW supports vector<4x4xf32>.
+/// Materialization is achieved by instantiating each occurrence of the leading
+/// dimension of vector<4x4x4xf32> into a vector<4x4xf32>.
+/// The program transformation that implements this instantiation is a
+/// multi-loop unroll-and-jam (it can be partial or full depending on the ratio
+/// of super-vector shape to HW-vector shape).
+///
+/// As a simple case, the following:
+/// ```mlir
+/// mlfunc @materialize(%M : index, %N : index, %O : index, %P : index) {
+/// %A = alloc (%M, %N, %O, %P) : memref<?x?x?x?xf32, 0>
+/// %f1 = constant splat<vector<4x4x4xf32>, 1.000000e+00> :
+/// vector<4x4x4xf32> for %i0 = 0 to %M step 4 {
+/// for %i1 = 0 to %N step 4 {
+/// for %i2 = 0 to %O {
+/// for %i3 = 0 to %P step 4 {
+/// vector_transfer_write %f1, %A, %i0, %i1, %i2, %i3
+/// {permutation_map: (d0, d1, d2, d3) -> (d3, d1, d0)} :
+/// vector<4x4x4xf32>, memref<?x?x?x?xf32, 0>,
+/// index, index, index, index
+/// }}}}
+/// return
+/// }
+/// ```
+///
+/// is instantiated by unroll-and-jam (just unroll in this case) into:
+///
+/// ```mlir
+/// mlfunc @materialize(%M : index, %N : index, %O : index, %P : index) {
+/// %A = alloc (%M, %N, %O, %P) : memref<?x?x?x?xf32, 0>
+/// %f1 = constant splat<vector<4x4xf32>, 1.000000e+00> : vector<4x4x4xf32>
+/// for %i0 = 0 to %arg0 step 4 {
+/// for %i1 = 0 to %arg1 step 4 {
+/// for %i2 = 0 to %arg2 {
+/// for %i3 = 0 to %arg3 step 4 {
+/// %1 = affine_apply (d0, d1, d2, d3) -> (d0, d1, d2, d3)
+/// (%i0, %i1, %i2, %i3)
+/// vector_transfer_write f1, %0, %1#0, %1#1, %1#2, %1#3
+/// {permutation_map: (d0, d1, d2, d3) -> (d1, d0)} :
+/// vector<4x4xf32>, memref<?x?x?x?xf32>,
+/// index, index, index, index
+/// %2 = affine_apply (d0, d1, d2, d3) -> (d0, d1, d2, d3 + 1)
+/// (%i0, %i1, %i2, %i3)
+/// vector_transfer_write {{.*}}, %0, %2#0, %2#1, %2#2, %2#3
+/// {permutation_map: (d0, d1, d2, d3) -> (d1, d0)} :
+/// vector<4x4xf32>, memref<?x?x?x?xf32>,
+/// index, index, index, index
+/// %3 = affine_apply (d0, d1, d2, d3) -> (d0, d1, d2, d3 + 2)
+/// (%i0, %i1, %i2, %i3)
+/// vector_transfer_write {{.*}}, %0, %3#0, %3#1, %3#2, %3#3
+/// {permutation_map: (d0, d1, d2, d3) -> (d1, d0)} :
+/// vector<4x4xf32>, memref<?x?x?x?xf32>,
+/// index, index, index, index
+/// %4 = affine_apply (d0, d1, d2, d3) -> (d0, d1, d2, d3 + 3)
+/// (%i0, %i1, %i2, %i3)
+/// vector_transfer_write {{.*}}, %0, %4#0, %4#1, %4#2, %4#3
+/// {permutation_map: (d0, d1, d2, d3) -> (d1, d0)} :
+/// vector<4x4xf32>, memref<?x?x?x?xf32>,
+/// index, index, index, index
+/// }}}}
+/// return
+/// }
+/// ```
using llvm::dbgs;
using llvm::DenseSet;
@@ -333,6 +400,58 @@
materializeAttributes(opStmt, superVectorType, hwVectorType));
}
+/// Computes the permutationMap required for a VectorTransferOp from the memref
+/// to the `hwVectorType`.
+/// This is achieved by returning the projection of the permutationMap along the
+/// dimensions of the super-vector type that remain in the hwVectorType.
+/// In particular, if a dimension is fully instantiated (i.e. unrolled) then it
+/// is projected out in the final result.
+template <typename VectorTransferOpTy>
+static AffineMap projectedPermutationMap(VectorTransferOpTy *transfer,
+ VectorType hwVectorType) {
+ static_assert(
+ std::is_same<VectorTransferOpTy, VectorTransferReadOp>::value ||
+ std::is_same<VectorTransferOpTy, VectorTransferWriteOp>::value,
+ "Must be called on a VectorTransferOp");
+ auto superVectorType = transfer->getVectorType();
+ auto optionalRatio = shapeRatio(superVectorType, hwVectorType);
+ assert(optionalRatio &&
+ (optionalRatio->size() == superVectorType.getShape().size()) &&
+ "Shape and ratio not of the same size");
+ unsigned dim = 0;
+ SmallVector<AffineExpr, 4> keep;
+ MLIRContext *context = transfer->getOperation()->getContext();
+ functional::zipApply(
+ [&dim, &keep, context](int shape, int ratio) {
+ assert(shape >= ratio && "shape dim must be greater than ratio dim");
+ if (shape != ratio) {
+ // HW vector is not full instantiated along this dim, keep it.
+ keep.push_back(getAffineDimExpr(dim, context));
+ }
+ ++dim;
+ },
+ superVectorType.getShape(), *optionalRatio);
+ auto projectionMap = AffineMap::get(optionalRatio->size(), 0, keep, {});
+ (void)projectionMap;
+ // No seemingly simple way to compose 2 AffineMap except going through SSA
+ // values... Punting for now and will resolve in the next CL.
+ //
+ // return projectionMap.compose(transfer->getPermutationMap());
+
+ // Still, we may need to drop a few dims to pass verification, so hack this in
+ // for now.
+ auto map = transfer->getPermutationMap();
+ auto exprs = map.getResults();
+ assert(exprs.size() >= keep.size());
+ unsigned diff = exprs.size() - keep.size();
+ SmallVector<AffineExpr, 4> projectedExprs(exprs.begin() + diff, exprs.end());
+ auto res = AffineMap::get(map.getNumInputs(), 0, projectedExprs, {});
+ LLVM_DEBUG(projectionMap.print(dbgs() << "\nProjectionMap: "));
+ LLVM_DEBUG(map.print(dbgs() << "\nOriginal: "));
+ LLVM_DEBUG(res.print(dbgs() << "\nTemporarily hacked projection: "));
+ return res;
+}
+
/// Creates an instantiated version of `read` for the instance of
/// `hwVectorInstance` when lowering from a super-vector type to
/// `hwVectorType`. `hwVectorInstance` represents one particular instance of
@@ -349,8 +468,7 @@
reindexAffineIndices(b, hwVectorType, hwVectorInstance, indices);
auto cloned = b->create<VectorTransferReadOp>(
read->getLoc(), hwVectorType, read->getMemRef(), affineIndices,
- makePermutationMap(read->getMemRefType(), hwVectorType),
- read->getPaddingValue());
+ projectedPermutationMap(read, hwVectorType), read->getPaddingValue());
return cast<OperationStmt>(cloned->getOperation());
}
@@ -371,7 +489,7 @@
auto cloned = b->create<VectorTransferWriteOp>(
write->getLoc(), substitute(write->getVector(), *substitutionsMap),
write->getMemRef(), affineIndices,
- makePermutationMap(write->getMemRefType(), hwVectorType));
+ projectedPermutationMap(write, hwVectorType));
return cast<OperationStmt>(cloned->getOperation());
}
diff --git a/lib/Transforms/Vectorize.cpp b/lib/Transforms/Vectorize.cpp
index e4822c2..1e41021 100644
--- a/lib/Transforms/Vectorize.cpp
+++ b/lib/Transforms/Vectorize.cpp
@@ -671,6 +671,18 @@
} // end anonymous namespace
+static void vectorizeLoopIfProfitable(ForStmt *loop, unsigned depthInPattern,
+ unsigned patternDepth,
+ VectorizationStrategy *strategy) {
+ assert(patternDepth > depthInPattern);
+ if (patternDepth - depthInPattern > strategy->vectorSizes.size()) {
+ // Don't vectorize this loop
+ return;
+ }
+ strategy->loopToVectorDim[loop] =
+ strategy->vectorSizes.size() - (patternDepth - depthInPattern);
+}
+
/// Implements a simple strawman strategy for vectorization.
/// Given a matched pattern `matches` of depth `patternDepth`, this strategy
/// greedily assigns the fastest varying dimension ** of the vector ** to the
@@ -696,17 +708,11 @@
if (fail) {
return fail;
}
- assert(patternDepth > depthInPattern);
- if (patternDepth - depthInPattern <= strategy->vectorSizes.size()) {
- strategy->loopToVectorDim[loop] =
- strategy->vectorSizes.size() - (patternDepth - depthInPattern);
- } else {
- // Don't vectorize
- strategy->loopToVectorDim[loop] = -1;
- }
+ vectorizeLoopIfProfitable(loop, depthInPattern, patternDepth, strategy);
}
return false;
}
+
///// end TODO(ntv): Hoist to a VectorizationStrategy.cpp when appropriate /////
namespace {
@@ -799,39 +805,6 @@
////// TODO(ntv): Hoist to a VectorizationMaterialize.cpp when appropriate. ////
-/// Creates a vector_transfer_read that loads a scalar MemRef into a
-/// super-vector register.
-///
-/// Usage:
-/// This vector_transfer_read op will be implemented as a PseudoOp for
-/// different backends. In its current form it is only used to load into a
-/// vector; where the vector may have any shape that is some multiple of the
-/// hardware-specific vector size used to implement the PseudoOp efficiently.
-/// This is used to implement "non-effecting padding" for early vectorization
-/// and allows higher-level passes in the codegen to not worry about
-/// hardware-specific implementation details.
-///
-/// TODO(ntv):
-/// 1. implement this end-to-end for some backend;
-/// 2. support operation-specific padding values to properly implement
-/// "non-effecting padding";
-/// 3. support input map for on-the-fly transpositions (point 1 above);
-/// 4. support broadcast map (point 5 above).
-///
-/// TODO(andydavis,bondhugula,ntv):
-/// 1. generalize to support padding semantics and offsets within vector type.
-static OperationStmt *
-createVectorTransferRead(OperationStmt *loadOp, VectorType vectorType,
- SSAValue *srcMemRef, ArrayRef<SSAValue *> srcIndices) {
- auto memRefType = srcMemRef->getType().cast<MemRefType>();
- MLFuncBuilder b(loadOp);
- // TODO(ntv): neutral for noneffective padding.
- auto transfer = b.create<VectorTransferReadOp>(
- loadOp->getLoc(), vectorType, srcMemRef, srcIndices,
- makePermutationMap(memRefType, vectorType));
- return cast<OperationStmt>(transfer->getOperation());
-}
-
/// Handles the vectorization of load and store MLIR operations.
///
/// LoadOp operations are the roots of the vectorizeOperations call. They are
@@ -863,10 +836,17 @@
// TODO(ntv): increase the expressiveness power of vector_transfer operations
// as needed by various targets.
if (opStmt->template isa<LoadOp>()) {
- auto *transfer = createVectorTransferRead(
- opStmt, vectorType, memoryOp->getMemRef(),
- map(makePtrDynCaster<SSAValue>(), memoryOp->getIndices()));
- state->registerReplacement(opStmt, transfer);
+ auto permutationMap =
+ makePermutationMap(opStmt, state->strategy->loopToVectorDim);
+ LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap: ");
+ LLVM_DEBUG(permutationMap.print(dbgs()));
+ MLFuncBuilder b(opStmt);
+ auto transfer = b.create<VectorTransferReadOp>(
+ opStmt->getLoc(), vectorType, memoryOp->getMemRef(),
+ map(makePtrDynCaster<SSAValue>(), memoryOp->getIndices()),
+ permutationMap);
+ state->registerReplacement(opStmt,
+ cast<OperationStmt>(transfer->getOperation()));
} else {
state->registerTerminator(opStmt);
}
@@ -943,16 +923,13 @@
// 2. This loop may have been omitted from vectorization for various reasons
// (e.g. due to the performance model or pattern depth > vector size).
- assert(state->strategy->loopToVectorDim.count(loop));
- assert(state->strategy->loopToVectorDim.find(loop) !=
- state->strategy->loopToVectorDim.end() &&
- "Key not found");
- int vectorDim = state->strategy->loopToVectorDim.lookup(loop);
- if (vectorDim < 0) {
+ auto it = state->strategy->loopToVectorDim.find(loop);
+ if (it == state->strategy->loopToVectorDim.end()) {
return false;
}
// 3. Actual post-order transformation.
+ auto vectorDim = it->second;
assert(vectorDim < state->strategy->vectorSizes.size() &&
"vector dim overflow");
// a. get actual vector size
@@ -1077,40 +1054,6 @@
return nullptr;
};
-/// Creates and returns a vector_transfer_write operation, which writes back a
-/// super-vector register into a scalar MemRef.
-///
-/// Usage:
-/// This vector_transfer_write op will be implemented as a PseudoOp for
-/// different backends. In its current form it is only used to store from a
-/// vector; where the vector may have any shape that is some multiple of
-/// the hardware-specific vector size used to implement the PseudoOp
-/// efficiently. This is used to implement "non-effecting padding" for early
-/// vectorization and allows higher-level passes in the codegen to not worry
-/// about hardware-specific implementation details.
-///
-/// TODO(ntv):
-/// 1. implement this end-to-end for some backend;
-/// 2. support write-back in the presence of races and ;
-/// 3. support input map for counterpart of broadcast (point 1 above);
-/// 4. support dstMap for writing back in non-contiguous memory regions
-/// (point 4 above).
-static OperationStmt *createVectorTransferWrite(OperationStmt *storeOp,
- VectorizationState *state) {
- auto store = storeOp->cast<StoreOp>();
- auto *memRef = store->getMemRef();
- auto memRefType = memRef->getType().cast<MemRefType>();
- auto *value = store->getValueToStore();
- auto *vectorValue = vectorizeOperand(value, storeOp, state);
- auto vectorType = vectorValue->getType().cast<VectorType>();
- auto indices = map(makePtrDynCaster<SSAValue>(), store->getIndices());
- MLFuncBuilder b(storeOp);
- auto transfer = b.create<VectorTransferWriteOp>(
- storeOp->getLoc(), vectorValue, memRef, indices,
- makePermutationMap(memRefType, vectorType));
- return cast<OperationStmt>(transfer->getOperation());
-}
-
/// Encodes OperationStmt-specific behavior for vectorization. In general we
/// assume that all operands of an op must be vectorized but this is not always
/// true. In the future, it would be nice to have a trait that describes how a
@@ -1121,31 +1064,41 @@
/// Maybe some Ops are not vectorizable or require some tricky logic, we cannot
/// do one-off logic here; ideally it would be TableGen'd.
static OperationStmt *vectorizeOneOperationStmt(MLFuncBuilder *b,
- OperationStmt *stmt,
+ OperationStmt *opStmt,
VectorizationState *state) {
// Sanity checks.
- assert(!stmt->isa<LoadOp>() &&
+ assert(!opStmt->isa<LoadOp>() &&
"all loads must have already been fully vectorized independently");
- assert(!stmt->isa<VectorTransferReadOp>() &&
+ assert(!opStmt->isa<VectorTransferReadOp>() &&
"vector_transfer_read cannot be further vectorized");
- assert(!stmt->isa<VectorTransferWriteOp>() &&
+ assert(!opStmt->isa<VectorTransferWriteOp>() &&
"vector_transfer_write cannot be further vectorized");
- if (stmt->isa<StoreOp>()) {
- auto *res = createVectorTransferWrite(stmt, state);
- LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ vectorized store: ");
- LLVM_DEBUG(res->print(dbgs()));
- // Terminators are erased on the spot.
- stmt->erase();
+ if (auto store = opStmt->dyn_cast<StoreOp>()) {
+ auto *memRef = store->getMemRef();
+ auto *value = store->getValueToStore();
+ auto *vectorValue = vectorizeOperand(value, opStmt, state);
+ auto indices = map(makePtrDynCaster<SSAValue>(), store->getIndices());
+ MLFuncBuilder b(opStmt);
+ auto permutationMap =
+ makePermutationMap(opStmt, state->strategy->loopToVectorDim);
+ LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap: ");
+ LLVM_DEBUG(permutationMap.print(dbgs()));
+ auto transfer = b.create<VectorTransferWriteOp>(
+ opStmt->getLoc(), vectorValue, memRef, indices, permutationMap);
+ auto *res = cast<OperationStmt>(transfer->getOperation());
+ LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ vectorized store: " << *res);
+ // "Terminators" (i.e. StoreOps) are erased on the spot.
+ opStmt->erase();
return res;
}
auto types = map([state](SSAValue *v) { return getVectorType(v, *state); },
- stmt->getResults());
- auto vectorizeOneOperand = [stmt, state](SSAValue *op) {
- return vectorizeOperand(op, stmt, state);
+ opStmt->getResults());
+ auto vectorizeOneOperand = [opStmt, state](SSAValue *op) {
+ return vectorizeOperand(op, opStmt, state);
};
- auto operands = map(vectorizeOneOperand, stmt->getOperands());
+ auto operands = map(vectorizeOneOperand, opStmt->getOperands());
// Check whether a single operand is null. If so, vectorization failed.
bool success = llvm::any_of(operands, [](SSAValue *op) { return op; });
if (!success) {
@@ -1159,8 +1112,9 @@
// TODO(ntv): Is it worth considering an OperationStmt.clone operation
// which changes the type so we can promote an OperationStmt with less
// boilerplate?
- return cast<OperationStmt>(b->createOperation(
- stmt->getLoc(), stmt->getName(), operands, types, stmt->getAttrs()));
+ return cast<OperationStmt>(b->createOperation(opStmt->getLoc(),
+ opStmt->getName(), operands,
+ types, opStmt->getAttrs()));
}
/// Iterates over the OperationStmt in the loop and rewrites them using their
@@ -1313,18 +1267,26 @@
LLVM_DEBUG(dbgs() << "\n******************************************");
LLVM_DEBUG(dbgs() << "\n[early-vect] new pattern on MLFunction\n");
LLVM_DEBUG(f->print(dbgs()));
+ unsigned patternDepth = pat.getDepth();
auto matches = pat.match(f);
- VectorizationStrategy strategy;
- // TODO(ntv): depending on profitability, elect to reduce the vector size.
- strategy.vectorSizes = clVirtualVectorSize;
- auto fail = analyzeProfitability(matches, 0, pat.getDepth(), &strategy);
- if (fail) {
- continue;
+ // Iterate over all the top-level matches and vectorize eagerly.
+ // This automatically prunes intersecting matches.
+ for (auto m : matches) {
+ VectorizationStrategy strategy;
+ // TODO(ntv): depending on profitability, elect to reduce the vector size.
+ strategy.vectorSizes = clVirtualVectorSize;
+ auto fail = analyzeProfitability(m.second, 1, patternDepth, &strategy);
+ if (fail) {
+ continue;
+ }
+ auto *loop = cast<ForStmt>(m.first);
+ vectorizeLoopIfProfitable(loop, 0, patternDepth, &strategy);
+ // TODO(ntv): if pattern does not apply, report it; alter the
+ // cost/benefit.
+ fail = vectorizeRootMatches(matches, &strategy);
+ assert(!fail && "top-level failure should not happen");
+ // TODO(ntv): some diagnostics.
}
- // TODO(ntv): if pattern does not apply, report it; alter the cost/benefit.
- fail = vectorizeRootMatches(matches, &strategy);
- assert(!fail && "top-level failure should not happen");
- // TODO(ntv): some diagnotics.
}
LLVM_DEBUG(dbgs() << "\n");
return PassResult::Success;
diff --git a/test/IR/invalid-ops.mlir b/test/IR/invalid-ops.mlir
index 03474a8..a566379 100644
--- a/test/IR/invalid-ops.mlir
+++ b/test/IR/invalid-ops.mlir
@@ -353,7 +353,7 @@
bb0(%arg0 : memref<?x?xf32>):
%c3 = constant 3 : index
%cst = constant 3.0 : f32
- // expected-error@+1 {{requires a permutation_map that is an actual permutation}}
+ // expected-error@+1 {{requires a projected permutation_map (at most one dim or the zero constant can appear in each result)}}
%0 = vector_transfer_read %arg0, %c3, %c3 {permutation_map: (d0, d1)->(d0 + d1)} : (memref<?x?xf32>, index, index) -> vector<128xf32>
}
@@ -363,7 +363,7 @@
bb0(%arg0 : memref<?x?xf32>):
%c3 = constant 3 : index
%cst = constant 3.0 : f32
- // expected-error@+1 {{requires a permutation_map that is an actual permutation}}
+ // expected-error@+1 {{requires a projected permutation_map (at most one dim or the zero constant can appear in each result)}}
%0 = vector_transfer_read %arg0, %c3, %c3 {permutation_map: (d0, d1)->(d0 + 1)} : (memref<?x?xf32>, index, index) -> vector<128xf32>
}
// -----
@@ -372,7 +372,7 @@
bb0(%arg0 : memref<?x?x?xf32>):
%c3 = constant 3 : index
%cst = constant 3.0 : f32
- // expected-error@+1 {{requires a permutation_map that is a full column-rank permutation}}
+ // expected-error@+1 {{requires a permutation_map that is a permutation (found one dim used more than once)}}
%0 = vector_transfer_read %arg0, %c3, %c3, %c3 {permutation_map: (d0, d1, d2)->(d0, d0)} : (memref<?x?x?xf32>, index, index, index) -> vector<3x7xf32>
}
@@ -442,7 +442,7 @@
bb0(%arg0 : memref<?x?xf32>):
%c3 = constant 3 : index
%cst = constant splat<vector<128 x f32>, 3.0> : vector<128 x f32>
- // expected-error@+1 {{requires a permutation_map that is an actual permutation}}
+ // expected-error@+1 {{requires a projected permutation_map (at most one dim or the zero constant can appear in each result)}}
vector_transfer_write %cst, %arg0, %c3, %c3 {permutation_map: (d0, d1)->(d0 + d1)} : vector<128xf32>, memref<?x?xf32>, index, index
}
@@ -452,7 +452,7 @@
bb0(%arg0 : memref<?x?xf32>):
%c3 = constant 3 : index
%cst = constant splat<vector<128 x f32>, 3.0> : vector<128 x f32>
- // expected-error@+1 {{requires a permutation_map that is an actual permutation}}
+ // expected-error@+1 {{requires a projected permutation_map (at most one dim or the zero constant can appear in each result)}}
vector_transfer_write %cst, %arg0, %c3, %c3 {permutation_map: (d0, d1)->(d0 + 1)} : vector<128xf32>, memref<?x?xf32>, index, index
}
// -----
@@ -461,7 +461,7 @@
bb0(%arg0 : memref<?x?x?xf32>):
%c3 = constant 3 : index
%cst = constant splat<vector<3 x 7 x f32>, 3.0> : vector<3 x 7 x f32>
- // expected-error@+1 {{requires a permutation_map that is a full column-rank permutation}}
+ // expected-error@+1 {{requires a permutation_map that is a permutation (found one dim used more than once)}}
vector_transfer_write %cst, %arg0, %c3, %c3, %c3 {permutation_map: (d0, d1, d2)->(d0, d0)} : vector<3x7xf32>, memref<?x?x?xf32>, index, index, index
}
diff --git a/test/Transforms/Vectorize/materialize_vectors_1d_to_1d.mlir b/test/Transforms/Vectorize/materialize_vectors_1d_to_1d.mlir
index 19e580e..85649cc 100644
--- a/test/Transforms/Vectorize/materialize_vectors_1d_to_1d.mlir
+++ b/test/Transforms/Vectorize/materialize_vectors_1d_to_1d.mlir
@@ -14,32 +14,82 @@
%C = alloc (%M, %N) : memref<?x?xf32, 0>
%f1 = constant 1.0 : f32
%f2 = constant 2.0 : f32
+ // 4x unroll (jammed by construction).
+ // CHECK: for %i0 = 0 to %arg0 {
+ // CHECK: for %i1 = 0 to %arg1 step 32 {
+ // CHECK: [[CST0:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: [[CST1:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: [[CST2:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: [[CST3:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: [[VAL0:%.*]] = affine_apply [[MAP0]]{{.*}}
+ // CHECK: vector_transfer_write [[CST0]], {{.*}}, [[VAL0]]#0, [[VAL0]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
+ // CHECK: [[VAL1:%.*]] = affine_apply [[MAP1]]{{.*}}
+ // CHECK: vector_transfer_write [[CST1]], {{.*}}, [[VAL1]]#0, [[VAL1]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
+ // CHECK: [[VAL2:%.*]] = affine_apply [[MAP2]]{{.*}}
+ // CHECK: vector_transfer_write [[CST2]], {{.*}}, [[VAL2]]#0, [[VAL2]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
+ // CHECK: [[VAL3:%.*]] = affine_apply [[MAP3]]{{.*}}
+ // CHECK: vector_transfer_write [[CST3]], {{.*}}, [[VAL3]]#0, [[VAL3]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
+ //
for %i0 = 0 to %M {
for %i1 = 0 to %N {
// non-scoped %f1
- // CHECK does 4x unrolling.
- // CHECK: [[CST0:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
- // CHECK: [[CST1:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
- // CHECK: [[CST2:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
- // CHECK: [[CST3:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
- // CHECK: [[VAL0:%.*]] = affine_apply [[MAP0]]{{.*}}
- // CHECK: vector_transfer_write [[CST0]], {{.*}}, [[VAL0]]#0, [[VAL0]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
- // CHECK: [[VAL1:%.*]] = affine_apply [[MAP1]]{{.*}}
- // CHECK: vector_transfer_write [[CST1]], {{.*}}, [[VAL1]]#0, [[VAL1]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
- // CHECK: [[VAL2:%.*]] = affine_apply [[MAP2]]{{.*}}
- // CHECK:vector_transfer_write [[CST2]], {{.*}}, [[VAL2]]#0, [[VAL2]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
- // CHECK: [[VAL3:%.*]] = affine_apply [[MAP3]]{{.*}}
- // CHECK:vector_transfer_write [[CST3]], {{.*}}, [[VAL3]]#0, [[VAL3]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
- //
store %f1, %A[%i0, %i1] : memref<?x?xf32, 0>
}
}
+ // 4x unroll (jammed by construction).
+ // CHECK: for %i2 = 0 to %arg0 {
+ // CHECK: for %i3 = 0 to %arg1 step 32 {
+ // CHECK: [[CST0:%.*]] = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: [[CST1:%.*]] = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: [[CST2:%.*]] = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: [[CST3:%.*]] = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: [[VAL0:%.*]] = affine_apply [[MAP0]]{{.*}}
+ // CHECK: vector_transfer_write [[CST0]], {{.*}}, [[VAL0]]#0, [[VAL0]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
+ // CHECK: [[VAL1:%.*]] = affine_apply [[MAP1]]{{.*}}
+ // CHECK: vector_transfer_write [[CST1]], {{.*}}, [[VAL1]]#0, [[VAL1]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
+ // CHECK: [[VAL2:%.*]] = affine_apply [[MAP2]]{{.*}}
+ // CHECK: vector_transfer_write [[CST2]], {{.*}}, [[VAL2]]#0, [[VAL2]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
+ // CHECK: [[VAL3:%.*]] = affine_apply [[MAP3]]{{.*}}
+ // CHECK: vector_transfer_write [[CST3]], {{.*}}, [[VAL3]]#0, [[VAL3]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
+ //
for %i2 = 0 to %M {
for %i3 = 0 to %N {
// non-scoped %f2
store %f2, %B[%i2, %i3] : memref<?x?xf32, 0>
}
}
+ // 4x unroll (jammed by construction).
+ // CHECK: for %i4 = 0 to %arg0 {
+ // CHECK: for %i5 = 0 to %arg1 step 32 {
+ // CHECK: %11 = affine_apply #map0(%i4, %i5)
+ // CHECK: %12 = vector_transfer_read %0, %11#0, %11#1 {permutation_map: #[[map_proj_d0d1_d1]]} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %13 = affine_apply #map2(%i4, %i5)
+ // CHECK: %14 = vector_transfer_read %0, %13#0, %13#1 {permutation_map: #[[map_proj_d0d1_d1]]} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %15 = affine_apply #map3(%i4, %i5)
+ // CHECK: %16 = vector_transfer_read %0, %15#0, %15#1 {permutation_map: #[[map_proj_d0d1_d1]]} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %17 = affine_apply #map4(%i4, %i5)
+ // CHECK: %18 = vector_transfer_read %0, %17#0, %17#1 {permutation_map: #[[map_proj_d0d1_d1]]} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %19 = affine_apply #map0(%i4, %i5)
+ // CHECK: %20 = vector_transfer_read %1, %19#0, %19#1 {permutation_map: #[[map_proj_d0d1_d1]]} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %21 = affine_apply #map2(%i4, %i5)
+ // CHECK: %22 = vector_transfer_read %1, %21#0, %21#1 {permutation_map: #[[map_proj_d0d1_d1]]} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %23 = affine_apply #map3(%i4, %i5)
+ // CHECK: %24 = vector_transfer_read %1, %23#0, %23#1 {permutation_map: #[[map_proj_d0d1_d1]]} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %25 = affine_apply #map4(%i4, %i5)
+ // CHECK: %26 = vector_transfer_read %1, %25#0, %25#1 {permutation_map: #[[map_proj_d0d1_d1]]} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %27 = addf %12, %20 : vector<8xf32>
+ // CHECK: %28 = addf %14, %22 : vector<8xf32>
+ // CHECK: %29 = addf %16, %24 : vector<8xf32>
+ // CHECK: %30 = addf %18, %26 : vector<8xf32>
+ // CHECK: %31 = affine_apply #map0(%i4, %i5)
+ // CHECK: vector_transfer_write %27, %2, %31#0, %31#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %32 = affine_apply #map2(%i4, %i5)
+ // CHECK: vector_transfer_write %28, %2, %32#0, %32#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %33 = affine_apply #map3(%i4, %i5)
+ // CHECK: vector_transfer_write %29, %2, %33#0, %33#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %34 = affine_apply #map4(%i4, %i5)
+ // CHECK: vector_transfer_write %30, %2, %34#0, %34#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>, memref<?x?xf32>, index, index
+ //
for %i4 = 0 to %M {
for %i5 = 0 to %N {
%a5 = load %A[%i4, %i5] : memref<?x?xf32, 0>
diff --git a/test/Transforms/Vectorize/materialize_vectors_2d_to_1d.mlir b/test/Transforms/Vectorize/materialize_vectors_2d_to_1d.mlir
index c9af214..afa2b41 100644
--- a/test/Transforms/Vectorize/materialize_vectors_2d_to_1d.mlir
+++ b/test/Transforms/Vectorize/materialize_vectors_2d_to_1d.mlir
@@ -16,25 +16,106 @@
%C = alloc (%M, %N) : memref<?x?xf32, 0>
%f1 = constant 1.0 : f32
%f2 = constant 2.0 : f32
+ // (3x2)x unroll (jammed by construction).
+ // CHECK: for %i0 = 0 to %arg0 step 3 {
+ // CHECK: for %i1 = 0 to %arg1 step 16 {
+ // CHECK: %cst_1 = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: %cst_2 = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: %cst_3 = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: %cst_4 = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: %cst_5 = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: %cst_6 = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
+ // CHECK: %3 = affine_apply #map0(%i0, %i1)
+ // CHECK: vector_transfer_write %cst_1, %0, %3#0, %3#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %4 = affine_apply #map2(%i0, %i1)
+ // CHECK: vector_transfer_write %cst_2, %0, %4#0, %4#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %5 = affine_apply #map3(%i0, %i1)
+ // CHECK: vector_transfer_write %cst_3, %0, %5#0, %5#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %6 = affine_apply #map4(%i0, %i1)
+ // CHECK: vector_transfer_write %cst_4, %0, %6#0, %6#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %7 = affine_apply #map5(%i0, %i1)
+ // CHECK: vector_transfer_write %cst_5, %0, %7#0, %7#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %8 = affine_apply #map6(%i0, %i1)
+ // CHECK: vector_transfer_write %cst_6, %0, %8#0, %8#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
for %i0 = 0 to %M {
for %i1 = 0 to %N {
// non-scoped %f1
store %f1, %A[%i0, %i1] : memref<?x?xf32, 0>
}
}
+ // (3x2)x unroll (jammed by construction).
+ // CHECK: for %i2 = 0 to %arg0 step 3 {
+ // CHECK: for %i3 = 0 to %arg1 step 16 {
+ // CHECK: %cst_7 = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: %cst_8 = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: %cst_9 = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: %cst_10 = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: %cst_11 = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: %cst_12 = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
+ // CHECK: %9 = affine_apply #map0(%i2, %i3)
+ // CHECK: vector_transfer_write %cst_7, %1, %9#0, %9#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %10 = affine_apply #map2(%i2, %i3)
+ // CHECK: vector_transfer_write %cst_8, %1, %10#0, %10#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %11 = affine_apply #map3(%i2, %i3)
+ // CHECK: vector_transfer_write %cst_9, %1, %11#0, %11#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %12 = affine_apply #map4(%i2, %i3)
+ // CHECK: vector_transfer_write %cst_10, %1, %12#0, %12#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %13 = affine_apply #map5(%i2, %i3)
+ // CHECK: vector_transfer_write %cst_11, %1, %13#0, %13#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %14 = affine_apply #map6(%i2, %i3)
+ // CHECK: vector_transfer_write %cst_12, %1, %14#0, %14#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
for %i2 = 0 to %M {
for %i3 = 0 to %N {
// non-scoped %f2
// CHECK does (3x4)x unrolling.
- // CHECK-COUNT-6: {{.*}} = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
- // CHECK: [[VAL0:%.*]] = affine_apply [[MAP0]]{{.*}}
- // CHECK: vector_transfer_write {{.*}}, [[VAL0]]#0, [[VAL0]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
- // ... 4 other interleaved affine_apply, vector_transfer_write
- // CHECK: [[VAL5:%.*]] = affine_apply [[MAP5]]{{.*}}
- // CHECK: vector_transfer_write {{.*}}, [[VAL5]]#0, [[VAL5]]#1 {permutation_map: #[[map_proj_d0d1_d1]]} : vector<8xf32>
store %f2, %B[%i2, %i3] : memref<?x?xf32, 0>
}
}
+ // (3x2)x unroll (jammed by construction).
+ // CHECK: for %i4 = 0 to %arg0 step 3 {
+ // CHECK: for %i5 = 0 to %arg1 step 16 {
+ // CHECK: %15 = affine_apply #map0(%i4, %i5)
+ // CHECK: %16 = vector_transfer_read %0, %15#0, %15#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %17 = affine_apply #map2(%i4, %i5)
+ // CHECK: %18 = vector_transfer_read %0, %17#0, %17#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %19 = affine_apply #map3(%i4, %i5)
+ // CHECK: %20 = vector_transfer_read %0, %19#0, %19#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %21 = affine_apply #map4(%i4, %i5)
+ // CHECK: %22 = vector_transfer_read %0, %21#0, %21#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %23 = affine_apply #map5(%i4, %i5)
+ // CHECK: %24 = vector_transfer_read %0, %23#0, %23#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %25 = affine_apply #map6(%i4, %i5)
+ // CHECK: %26 = vector_transfer_read %0, %25#0, %25#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %27 = affine_apply #map0(%i4, %i5)
+ // CHECK: %28 = vector_transfer_read %1, %27#0, %27#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %29 = affine_apply #map2(%i4, %i5)
+ // CHECK: %30 = vector_transfer_read %1, %29#0, %29#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %31 = affine_apply #map3(%i4, %i5)
+ // CHECK: %32 = vector_transfer_read %1, %31#0, %31#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %33 = affine_apply #map4(%i4, %i5)
+ // CHECK: %34 = vector_transfer_read %1, %33#0, %33#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %35 = affine_apply #map5(%i4, %i5)
+ // CHECK: %36 = vector_transfer_read %1, %35#0, %35#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %37 = affine_apply #map6(%i4, %i5)
+ // CHECK: %38 = vector_transfer_read %1, %37#0, %37#1 {permutation_map: #map1} : (memref<?x?xf32>, index, index) -> vector<8xf32>
+ // CHECK: %39 = addf %16, %28 : vector<8xf32>
+ // CHECK: %40 = addf %18, %30 : vector<8xf32>
+ // CHECK: %41 = addf %20, %32 : vector<8xf32>
+ // CHECK: %42 = addf %22, %34 : vector<8xf32>
+ // CHECK: %43 = addf %24, %36 : vector<8xf32>
+ // CHECK: %44 = addf %26, %38 : vector<8xf32>
+ // CHECK: %45 = affine_apply #map0(%i4, %i5)
+ // CHECK: vector_transfer_write %39, %2, %45#0, %45#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %46 = affine_apply #map2(%i4, %i5)
+ // CHECK: vector_transfer_write %40, %2, %46#0, %46#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %47 = affine_apply #map3(%i4, %i5)
+ // CHECK: vector_transfer_write %41, %2, %47#0, %47#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %48 = affine_apply #map4(%i4, %i5)
+ // CHECK: vector_transfer_write %42, %2, %48#0, %48#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %49 = affine_apply #map5(%i4, %i5)
+ // CHECK: vector_transfer_write %43, %2, %49#0, %49#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
+ // CHECK: %50 = affine_apply #map6(%i4, %i5)
+ // CHECK: vector_transfer_write %44, %2, %50#0, %50#1 {permutation_map: #map1} : vector<8xf32>, memref<?x?xf32>, index, index
for %i4 = 0 to %M {
for %i5 = 0 to %N {
%a5 = load %A[%i4, %i5] : memref<?x?xf32, 0>
diff --git a/test/Transforms/Vectorize/materialize_vectors_2d_to_2d.mlir b/test/Transforms/Vectorize/materialize_vectors_2d_to_2d.mlir
index b7ee8e6..b1a9cc7 100644
--- a/test/Transforms/Vectorize/materialize_vectors_2d_to_2d.mlir
+++ b/test/Transforms/Vectorize/materialize_vectors_2d_to_2d.mlir
@@ -11,35 +11,58 @@
%C = alloc (%M, %N) : memref<?x?xf32, 0>
%f1 = constant 1.0 : f32
%f2 = constant 2.0 : f32
+ // 2x unroll (jammed by construction).
+ // CHECK: for %i0 = 0 to %arg0 step 3 {
+ // CHECK: for %i1 = 0 to %arg1 step 32 {
+ // CHECK: %cst_1 = constant splat<vector<3x16xf32>, 1.000000e+00> : vector<3x16xf32>
+ // CHECK: %cst_2 = constant splat<vector<3x16xf32>, 1.000000e+00> : vector<3x16xf32>
+ // CHECK: %3 = affine_apply #map0(%i0, %i1)
+ // CHECK: vector_transfer_write %cst_1, %0, %3#0, %3#1 {permutation_map: #map0} : vector<3x16xf32>, memref<?x?xf32>, index, index
+ // CHECK: %4 = affine_apply #map1(%i0, %i1)
+ // CHECK: vector_transfer_write %cst_2, %0, %4#0, %4#1 {permutation_map: #map0} : vector<3x16xf32>, memref<?x?xf32>, index, index
+ //
for %i0 = 0 to %M {
for %i1 = 0 to %N {
// non-scoped %f1
store %f1, %A[%i0, %i1] : memref<?x?xf32, 0>
}
}
+ // 2x unroll (jammed by construction).
+ // CHECK: for %i2 = 0 to %arg0 step 3 {
+ // CHECK: for %i3 = 0 to %arg1 step 32 {
+ // CHECK: %cst_3 = constant splat<vector<3x16xf32>, 2.000000e+00> : vector<3x16xf32>
+ // CHECK: %cst_4 = constant splat<vector<3x16xf32>, 2.000000e+00> : vector<3x16xf32>
+ // CHECK: %5 = affine_apply #map0(%i2, %i3)
+ // CHECK: vector_transfer_write %cst_3, %1, %5#0, %5#1 {permutation_map: #map0} : vector<3x16xf32>, memref<?x?xf32>, index, index
+ // CHECK: %6 = affine_apply #map1(%i2, %i3)
+ // CHECK: vector_transfer_write %cst_4, %1, %6#0, %6#1 {permutation_map: #map0} : vector<3x16xf32>, memref<?x?xf32>, index, index
+ //
for %i2 = 0 to %M {
for %i3 = 0 to %N {
// non-scoped %f2
store %f2, %B[%i2, %i3] : memref<?x?xf32, 0>
}
}
+ // 2x unroll (jammed by construction).
+ // CHECK: for %i4 = 0 to %arg0 step 3 {
+ // CHECK: for %i5 = 0 to %arg1 step 32 {
+ // CHECK: %7 = affine_apply #map0(%i4, %i5)
+ // CHECK: %8 = vector_transfer_read %0, %7#0, %7#1 {permutation_map: #map0} : (memref<?x?xf32>, index, index) -> vector<3x16xf32>
+ // CHECK: %9 = affine_apply #map1(%i4, %i5)
+ // CHECK: %10 = vector_transfer_read %0, %9#0, %9#1 {permutation_map: #map0} : (memref<?x?xf32>, index, index) -> vector<3x16xf32>
+ // CHECK: %11 = affine_apply #map0(%i4, %i5)
+ // CHECK: %12 = vector_transfer_read %1, %11#0, %11#1 {permutation_map: #map0} : (memref<?x?xf32>, index, index) -> vector<3x16xf32>
+ // CHECK: %13 = affine_apply #map1(%i4, %i5)
+ // CHECK: %14 = vector_transfer_read %1, %13#0, %13#1 {permutation_map: #map0} : (memref<?x?xf32>, index, index) -> vector<3x16xf32>
+ // CHECK: %15 = addf %8, %12 : vector<3x16xf32>
+ // CHECK: %16 = addf %10, %14 : vector<3x16xf32>
+ // CHECK: %17 = affine_apply #map0(%i4, %i5)
+ // CHECK: vector_transfer_write %15, %2, %17#0, %17#1 {permutation_map: #map0} : vector<3x16xf32>, memref<?x?xf32>, index, index
+ // CHECK: %18 = affine_apply #map1(%i4, %i5)
+ // CHECK: vector_transfer_write %16, %2, %18#0, %18#1 {permutation_map: #map0} : vector<3x16xf32>, memref<?x?xf32>, index, index
+ //
for %i4 = 0 to %M {
for %i5 = 0 to %N {
- // CHECK: %7 = affine_apply #map0(%i4, %i5)
- // CHECK: %8 = vector_transfer_read %0, %7#0, %7#1 {permutation_map: #[[map_proj_d0d1_d0d1]]} : (memref<?x?xf32>, index, index) -> vector<3x16xf32>
- // CHECK: %9 = affine_apply #map1(%i4, %i5)
- // CHECK: %10 = vector_transfer_read %0, %9#0, %9#1 {permutation_map: #[[map_proj_d0d1_d0d1]]} : (memref<?x?xf32>, index, index) -> vector<3x16xf32>
- // CHECK: %11 = affine_apply #map0(%i4, %i5)
- // CHECK: %12 = vector_transfer_read %1, %11#0, %11#1 {permutation_map: #[[map_proj_d0d1_d0d1]]} : (memref<?x?xf32>, index, index) -> vector<3x16xf32>
- // CHECK: %13 = affine_apply #map1(%i4, %i5)
- // CHECK: %14 = vector_transfer_read %1, %13#0, %13#1 {permutation_map: #[[map_proj_d0d1_d0d1]]} : (memref<?x?xf32>, index, index) -> vector<3x16xf32>
- // CHECK: %15 = addf %8, %12 : vector<3x16xf32>
- // CHECK: %16 = addf %10, %14 : vector<3x16xf32>
- // CHECK: %17 = affine_apply #map0(%i4, %i5)
- // CHECK: vector_transfer_write %15, %2, %17#0, %17#1 {permutation_map: #[[map_proj_d0d1_d0d1]]} : vector<3x16xf32>, memref<?x?xf32>, index, index
- // CHECK: %18 = affine_apply #map1(%i4, %i5)
- // CHECK: vector_transfer_write %16, %2, %18#0, %18#1 {permutation_map: #[[map_proj_d0d1_d0d1]]} : vector<3x16xf32>, memref<?x?xf32>, index, index
- //
%a5 = load %A[%i4, %i5] : memref<?x?xf32, 0>
%b5 = load %B[%i4, %i5] : memref<?x?xf32, 0>
%s5 = addf %a5, %b5 : f32
diff --git a/test/Transforms/Vectorize/vectorize_1d.mlir b/test/Transforms/Vectorize/vectorize_1d.mlir
index da3cedc..1c37563 100644
--- a/test/Transforms/Vectorize/vectorize_1d.mlir
+++ b/test/Transforms/Vectorize/vectorize_1d.mlir
@@ -1,6 +1,7 @@
// RUN: mlir-opt %s -vectorize -virtual-vector-size 128 --test-fastest-varying=0 | FileCheck %s
// Permutation maps used in vectorization.
+// CHECK: #[[map_proj_d0d1_0:map[0-9]+]] = (d0, d1) -> (0)
// CHECK: #[[map_proj_d0d1_d1:map[0-9]+]] = (d0, d1) -> (d1)
#map0 = (d0) -> (d0)
@@ -24,24 +25,24 @@
%N = dim %A, 1 : memref<?x?xf32>
%P = dim %B, 2 : memref<?x?x?xf32>
%cst0 = constant 0 : index
-// CHECK:for [[IV0:%[a-zA-Z0-9]+]] = 0 to [[ARG_M]] step 128
-// CHECK-NEXT: {{.*}} = vector_transfer_read %arg0, [[C0]], [[C0]] {permutation_map: #[[map_proj_d0d1_d1]]} : (memref<?x?xf32>, index, index) -> vector<128xf32>
-// For this simple loop, the current transformation generates:
-// for %i0 = 0 to %0 step 128 {
-// %3 = vector_transfer_read %arg0, %c0_0, %c0_0 : (memref<?x?xf32>, index, index) -> vector<128xf32>
-// }
+//
+// CHECK: for {{.*}} step 128
+// CHECK-NEXT: {{.*}} = vector_transfer_read %arg0, [[C0]], [[C0]] {permutation_map: #[[map_proj_d0d1_0]]} : (memref<?x?xf32>, index, index) -> vector<128xf32>
for %i0 = 0 to %M { // vectorized due to scalar -> vector
%a0 = load %A[%cst0, %cst0] : memref<?x?xf32>
}
+//
// CHECK:for {{.*}} [[ARG_M]] {
for %i1 = 0 to %M { // not vectorized
%a1 = load %A[%i1, %i1] : memref<?x?xf32>
}
+//
// CHECK: for %i{{[0-9]*}} = 0 to [[ARG_M]] {
for %i2 = 0 to %M { // not vectorized, would vectorize with --test-fastest-varying=1
%r2 = affine_apply (d0) -> (d0) (%i2)
%a2 = load %A[%r2#0, %cst0] : memref<?x?xf32>
}
+//
// CHECK:for [[IV3:%[a-zA-Z0-9]+]] = 0 to [[ARG_M]] step 128
// CHECK-NEXT: [[APP3:%[a-zA-Z0-9]+]] = affine_apply {{.*}}[[IV3]]
// CHECK-NEXT: {{.*}} = vector_transfer_read %arg0, [[C0]], [[APP3]] {permutation_map: #[[map_proj_d0d1_d1]]} : {{.*}} -> vector<128xf32>
@@ -49,6 +50,7 @@
%r3 = affine_apply (d0) -> (d0) (%i3)
%a3 = load %A[%cst0, %r3#0] : memref<?x?xf32>
}
+//
// CHECK:for [[IV4:%[i0-9]+]] = 0 to [[ARG_M]] step 128 {
// CHECK-NEXT: for [[IV5:%[i0-9]*]] = 0 to [[ARG_N]] {
// CHECK-NEXT: [[APP5:%[0-9]+]] = affine_apply {{.*}}([[IV4]], [[IV5]])
@@ -59,6 +61,7 @@
%a5 = load %A[%r5#0, %r5#1] : memref<?x?xf32>
}
}
+//
// CHECK: for [[IV6:%[i0-9]*]] = 0 to [[ARG_M]] {
// CHECK-NEXT: for [[IV7:%[i0-9]*]] = 0 to [[ARG_N]] {
for %i6 = 0 to %M { // not vectorized, would vectorize with --test-fastest-varying=1
@@ -67,6 +70,7 @@
%a7 = load %A[%r7#0, %r7#1] : memref<?x?xf32>
}
}
+//
// CHECK:for [[IV8:%[i0-9]+]] = 0 to [[ARG_M]] step 128
// CHECK-NEXT: for [[IV9:%[i0-9]*]] = 0 to [[ARG_N]] {
// CHECK-NEXT: [[APP9:%[0-9]+]] = affine_apply {{.*}}([[IV8]], [[IV9]])
@@ -77,6 +81,7 @@
%a9 = load %A[%r9#0, %r9#1] : memref<?x?xf32>
}
}
+//
// CHECK: for [[IV10:%[i0-9]*]] = 0 to %{{[0-9]*}} {
// CHECK: for [[IV11:%[i0-9]*]] = 0 to %{{[0-9]*}} {
for %i10 = 0 to %M { // not vectorized, need per load transposes
@@ -87,6 +92,7 @@
store %a11, %A[%r12#0, %r12#1] : memref<?x?xf32>
}
}
+//
// CHECK: for [[IV12:%[i0-9]*]] = 0 to %{{[0-9]*}} {
// CHECK: for [[IV13:%[i0-9]*]] = 0 to %{{[0-9]*}} {
// CHECK: for [[IV14:%[i0-9]+]] = 0 to [[ARG_P]] step 128
@@ -98,20 +104,23 @@
}
}
}
+//
// CHECK: for %i{{[0-9]*}} = 0 to %{{[0-9]*}} {
for %i15 = 0 to %M { // not vectorized due to condition below
if #set0(%i15) {
%a15 = load %A[%cst0, %cst0] : memref<?x?xf32>
}
}
+//
// CHECK: for %i{{[0-9]*}} = 0 to %{{[0-9]*}} {
for %i16 = 0 to %M { // not vectorized, can't vectorize a vector load
%a16 = alloc(%M) : memref<?xvector<2xf32>>
%l16 = load %a16[%i16] : memref<?xvector<2xf32>>
}
+//
// CHECK: for %i{{[0-9]*}} = 0 to %{{[0-9]*}} {
// CHECK: for [[IV18:%[a-zA-Z0-9]+]] = 0 to [[ARG_M]] step 128
-// CHECK: {{.*}} = vector_transfer_read %arg0, [[C0]], [[C0]] {permutation_map: #[[map_proj_d0d1_d1]]} : {{.*}} -> vector<128xf32>
+// CHECK: {{.*}} = vector_transfer_read %arg0, [[C0]], [[C0]] {permutation_map: #[[map_proj_d0d1_0]]} : {{.*}} -> vector<128xf32>
for %i17 = 0 to %M { // not vectorized, the 1-D pattern that matched %i18 in DFS post-order prevents vectorizing %i17
for %i18 = 0 to %M { // vectorized due to scalar -> vector
%a18 = load %A[%cst0, %cst0] : memref<?x?xf32>
diff --git a/test/Transforms/Vectorize/vectorize_outer_loop_2d.mlir b/test/Transforms/Vectorize/vectorize_outer_loop_2d.mlir
index 807f304..486cb2a 100644
--- a/test/Transforms/Vectorize/vectorize_outer_loop_2d.mlir
+++ b/test/Transforms/Vectorize/vectorize_outer_loop_2d.mlir
@@ -1,19 +1,16 @@
// RUN: mlir-opt %s -vectorize -virtual-vector-size 32 -virtual-vector-size 256 --test-fastest-varying=2 --test-fastest-varying=0 | FileCheck %s
+
// Permutation maps used in vectorization.
-// CHECK: #[[map_proj_d0d1d2_d1d2:map[0-9]+]] = (d0, d1, d2) -> (d1, d2)
+// CHECK: #[[map_proj_d0d1d2_d0d2:map[0-9]+]] = (d0, d1, d2) -> (d0, d2)
mlfunc @vec2d(%A : memref<?x?x?xf32>) {
%M = dim %A, 0 : memref<?x?x?xf32>
%N = dim %A, 1 : memref<?x?x?xf32>
%P = dim %A, 2 : memref<?x?x?xf32>
- // CHECK: for {{.*}} = 0 to %0 step 32
- // CHECK: for {{.*}} = 0 to %1 {
- // CHECK: for {{.*}} = 0 to %2 step 256
- // Expected:
- // for %i0 = 0 to %0 step 32 {
- // for %i1 = 0 to %1 {
- // for %i2 = 0 to %2 step 256 {
- // %3 = "vector_transfer_read"(%arg0, %i0, %i1, %i2) : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
+ // CHECK: for %i0 = 0 to %0 step 32
+ // CHECK: for %i1 = 0 to %1 {
+ // CHECK: for %i2 = 0 to %2 step 256
+ // CHECK: {{.*}} = vector_transfer_read %arg0, %i0, %i1, %i2 {permutation_map: #[[map_proj_d0d1d2_d0d2]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
for %i0 = 0 to %M {
for %i1 = 0 to %N {
for %i2 = 0 to %P {
diff --git a/test/Transforms/Vectorize/vectorize_outer_loop_transpose_2d.mlir b/test/Transforms/Vectorize/vectorize_outer_loop_transpose_2d.mlir
index 420569f..8c05d28 100644
--- a/test/Transforms/Vectorize/vectorize_outer_loop_transpose_2d.mlir
+++ b/test/Transforms/Vectorize/vectorize_outer_loop_transpose_2d.mlir
@@ -1,7 +1,7 @@
// RUN: mlir-opt %s -vectorize -virtual-vector-size 32 -virtual-vector-size 256 --test-fastest-varying=0 --test-fastest-varying=2 | FileCheck %s
// Permutation maps used in vectorization.
-// CHECK: #[[map_proj_d0d1d2_d1d2:map[0-9]+]] = (d0, d1, d2) -> (d1, d2)
+// CHECK: #[[map_proj_d0d1d2_d2d0:map[0-9]+]] = (d0, d1, d2) -> (d2, d0)
mlfunc @vec2d(%A : memref<?x?x?xf32>) {
%M = dim %A, 0 : memref<?x?x?xf32>
@@ -19,14 +19,10 @@
}
}
}
- // CHECK: for {{.*}} = 0 to %0 step 32
- // CHECK: for {{.*}} = 0 to %1 step 256
- // CHECK: for {{.*}} = 0 to %2 {
- // Expected:
- // for %i3 = 0 to %0 step 32 {
- // for %i4 = 0 to %1 step 256 {
- // for %i5 = 0 to %2 {
- // %4 = "vector_transfer_read"(%arg0, %i4, %i5, %i3, %4) : (memref<?x?x?xf32>, index, index) -> vector<32x256xf32>
+ // CHECK: for %i3 = 0 to %0 step 32
+ // CHECK: for %i4 = 0 to %1 step 256
+ // CHECK: for %i5 = 0 to %2 {
+ // CHECK: {{.*}} = vector_transfer_read %arg0, %i4, %i5, %i3 {permutation_map: #[[map_proj_d0d1d2_d2d0]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
for %i3 = 0 to %M {
for %i4 = 0 to %N {
for %i5 = 0 to %P {
@@ -44,12 +40,12 @@
// CHECK: for %i0 = 0 to %0 step 32 {
// CHECK: for %i1 = 0 to %1 {
// CHECK: for %i2 = 0 to %2 step 256 {
- // CHECK: %3 = vector_transfer_read %arg0, %i2, %i1, %i0 {permutation_map: #[[map_proj_d0d1d2_d1d2]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
+ // CHECK: %3 = vector_transfer_read %arg0, %i2, %i1, %i0 {permutation_map: #[[map_proj_d0d1d2_d2d0]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
// CHECK: for %i3 = 0 to %1 step 256 {
// CHECK: for %i4 = 0 to %2 {
- // CHECK: %4 = vector_transfer_read %arg0, %i3, %i4, %i0 {permutation_map: #[[map_proj_d0d1d2_d1d2]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
+ // CHECK: %4 = vector_transfer_read %arg0, %i3, %i4, %i0 {permutation_map: #[[map_proj_d0d1d2_d2d0]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
// CHECK: for %i5 = 0 to %2 {
- // CHECK: %5 = vector_transfer_read %arg0, %i3, %i5, %i0 {permutation_map: #[[map_proj_d0d1d2_d1d2]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
+ // CHECK: %5 = vector_transfer_read %arg0, %i3, %i5, %i0 {permutation_map: #[[map_proj_d0d1d2_d2d0]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
for %i0 = 0 to %0 {
for %i1 = 0 to %1 {
for %i2 = 0 to %2 {
diff --git a/test/Transforms/Vectorize/vectorize_transpose_2d.mlir b/test/Transforms/Vectorize/vectorize_transpose_2d.mlir
index cbc84b7..bc82a0e 100644
--- a/test/Transforms/Vectorize/vectorize_transpose_2d.mlir
+++ b/test/Transforms/Vectorize/vectorize_transpose_2d.mlir
@@ -1,7 +1,7 @@
// RUN: mlir-opt %s -vectorize -virtual-vector-size 32 -virtual-vector-size 256 --test-fastest-varying=0 --test-fastest-varying=1 | FileCheck %s
// Permutation maps used in vectorization.
-// CHECK: #[[map_proj_d0d1d2_d1d2:map[0-9]+]] = (d0, d1, d2) -> (d1, d2)
+// CHECK-DAG: #[[map_proj_d0d1d2_d2d1:map[0-9]+]] = (d0, d1, d2) -> (d2, d1)
mlfunc @vec2d(%A : memref<?x?x?xf32>) {
%M = dim %A, 0 : memref<?x?x?xf32>
@@ -19,14 +19,10 @@
}
}
}
- // CHECK: for {{.*}} = 0 to %0 step 32
- // CHECK: for {{.*}} = 0 to %1 {
- // CHECK: for {{.*}} = 0 to %2 step 256
- // For the case: --test-fastest-varying=0 --test-fastest-varying=1:
- // for %i3 = 0 to %0 step 32 {
- // for %i4 = 0 to %1 {
- // for %i5 = 0 to %2 step 256 {
- // %4 = "vector_transfer_read"(%arg0, %i4, %i5, %i3, %4) : (memref<?x?x?xf32>, index, index) -> vector<32x256xf32>
+ // CHECK: for %i3 = 0 to %0 step 32
+ // CHECK: for %i4 = 0 to %1 {
+ // CHECK: for %i5 = 0 to %2 step 256
+ // CHECK: {{.*}} = vector_transfer_read %arg0, %i4, %i5, %i3 {permutation_map: #[[map_proj_d0d1d2_d2d1]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
for %i3 = 0 to %M {
for %i4 = 0 to %N {
for %i5 = 0 to %P {
@@ -44,12 +40,12 @@
// CHECK: for %i0 = 0 to %0 step 32 {
// CHECK: for %i1 = 0 to %1 step 256 {
// CHECK: for %i2 = 0 to %2 {
- // CHECK: %3 = vector_transfer_read %arg0, %i2, %i1, %i0 {permutation_map: #[[map_proj_d0d1d2_d1d2]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
+ // CHECK: %3 = vector_transfer_read %arg0, %i2, %i1, %i0 {permutation_map: #[[map_proj_d0d1d2_d2d1]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
// CHECK: for %i3 = 0 to %1 {
// CHECK: for %i4 = 0 to %2 step 256 {
- // CHECK: %4 = vector_transfer_read %arg0, %i3, %i4, %i0 {permutation_map: #[[map_proj_d0d1d2_d1d2]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
+ // CHECK: %4 = vector_transfer_read %arg0, %i3, %i4, %i0 {permutation_map: #[[map_proj_d0d1d2_d2d1]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
// CHECK: for %i5 = 0 to %2 step 256 {
- // CHECK: %5 = vector_transfer_read %arg0, %i3, %i5, %i0 {permutation_map: #[[map_proj_d0d1d2_d1d2]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
+ // CHECK: %5 = vector_transfer_read %arg0, %i3, %i5, %i0 {permutation_map: #[[map_proj_d0d1d2_d2d1]]} : (memref<?x?x?xf32>, index, index, index) -> vector<32x256xf32>
for %i0 = 0 to %0 {
for %i1 = 0 to %1 {
for %i2 = 0 to %2 {