[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/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,