Skip to content

[MLIR][Affine] Normalize memref.alloc ops with non trivial layout map #129875

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[MLIR][Affine] Normalize memref.alloc ops with non trivial layout map #129875

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions mlir/lib/Analysis/FlatLinearValueConstraints.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,9 +118,11 @@ struct SemiAffineExprFlattener : public AffineExprFlattener {
// with a positive value." (see AffineExprKind in AffineExpr.h). If this
// assumption does not hold constraints (added above) are a contradiction.

return success();
} else if (localExpr.getKind() == AffineExprKind::Mul) {
(void)localVarCst.appendVar(VarKind::Local);
return success();
}

// TODO: Support other semi-affine expressions.
return failure();
}
Expand Down Expand Up @@ -163,7 +165,6 @@ getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,

return success();
};

if (addConservativeSemiAffineBounds) {
SemiAffineExprFlattener flattener(numDims, numSymbols);
return flattenExprs(flattener);
Expand Down Expand Up @@ -229,7 +230,8 @@ LogicalResult FlatLinearConstraints::composeMatchingMap(AffineMap other) {
assert(other.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");

std::vector<SmallVector<int64_t, 8>> flatExprs;
if (failed(flattenAlignedMapAndMergeLocals(other, &flatExprs)))
if (failed(flattenAlignedMapAndMergeLocals(
other, &flatExprs, /*addConservativeSemiAffineBounds=*/true)))
return failure();
assert(flatExprs.size() == other.getNumResults());

Expand Down Expand Up @@ -796,8 +798,6 @@ LogicalResult FlatLinearConstraints::flattenAlignedMapAndMergeLocals(
<< "composition unimplemented for semi-affine maps\n");
return failure();
}

// Add localCst information.
if (localCst.getNumLocalVars() > 0) {
unsigned numLocalVars = getNumLocalVars();
// Insert local dims of localCst at the beginning.
Expand Down
56 changes: 41 additions & 15 deletions mlir/lib/Dialect/Affine/Utils/Utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1786,7 +1786,6 @@ static void createNewDynamicSizes(MemRefType oldMemRefType,
}
}

// TODO: Currently works for static memrefs with a single layout map.
template <typename AllocLikeOp>
LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
MemRefType memrefType = allocOp->getType();
Expand All @@ -1799,7 +1798,6 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
// Either memrefType already had an identity map or the map couldn't be
// transformed to an identity map.
return failure();

Value oldMemRef = allocOp->getResult();

SmallVector<Value, 4> symbolOperands(allocOp->getSymbolOperands());
Expand All @@ -1819,8 +1817,40 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType, newDynamicSizes,
allocOp->getAlignmentAttr());
} else {
newAlloc = b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType,
allocOp->getAlignmentAttr());
mlir::ValueRange dynamicSizes = allocOp->getDynamicSizes();
Copy link
Contributor

Choose a reason for hiding this comment

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

Need a comment as to what this else block is meant to handle.

mlir::ValueRange symbolOperands = allocOp->getSymbolOperands();
ArrayRef<int64_t> newShape = newMemRefType.getShape();
ArrayRef<int64_t> oldShape = memrefType.getShape();
SmallVector<Value> mapOperands(oldShape.size() + symbolOperands.size());
SmallVector<Value> dimensionOperands;
unsigned dimId = 0, symId = 0;
// Collect all the map operands of `allocOp` (both dynamic sizes and symbol
// operands), which will help us to compute the dynamic sizes of the new
// alloc op we are going to create.
for (unsigned i = 0, e = oldShape.size(); i < e; i++) {
if (oldShape[i] == ShapedType::kDynamic)
Copy link
Contributor

Choose a reason for hiding this comment

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

Use isDynamicDim

mapOperands[i] = dynamicSizes[dimId++];
else
mapOperands[i] =
b.create<arith::ConstantIndexOp>(allocOp->getLoc(), oldShape[i]);
}
for (unsigned i = oldShape.size(), e = mapOperands.size(); i < e; i++)
mapOperands[i] = symbolOperands[symId++];
// Compute the dynamic sizes operands for the new alloc op. If `newShape` is
// dynamic along a dimension, compute its shape using the layout map and
// dynamic sizes and symbol operands of the old `allocOp`.
for (unsigned i = 0, e = newShape.size(); i < e; i++) {
if (newShape[i] != ShapedType::kDynamic)
continue;
AffineExpr resExpr = layoutMap.getResult(i);
auto resMap = AffineMap::get(layoutMap.getNumDims(),
layoutMap.getNumSymbols(), resExpr);
dimensionOperands.push_back(
b.create<AffineApplyOp>(allocOp->getLoc(), resMap, mapOperands));
}
newAlloc =
b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType,
dimensionOperands, allocOp->getAlignmentAttr());
}
// Replace all uses of the old memref.
if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
Expand Down Expand Up @@ -1868,11 +1898,8 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {

// Normalize only static memrefs and dynamic memrefs with a tiled-layout map
// for now.
// TODO: Normalize the other types of dynamic memrefs.
SmallVector<std::tuple<AffineExpr, unsigned, unsigned>> tileSizePos;
(void)getTileSizePos(layoutMap, tileSizePos);
if (memrefType.getNumDynamicDims() > 0 && tileSizePos.empty())
return memrefType;

// We have a single map that is not an identity map. Create a new memref
// with the right shape and an identity layout map.
Expand All @@ -1894,7 +1921,6 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
unsigned newRank = layoutMap.getNumResults();
if (failed(fac.composeMatchingMap(layoutMap)))
return memrefType;
// TODO: Handle semi-affine maps.
// Project out the old data dimensions.
fac.projectOut(newRank, fac.getNumVars() - newRank - fac.getNumLocalVars());
SmallVector<int64_t, 4> newShape(newRank);
Expand All @@ -1910,14 +1936,14 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
// For a static memref and an affine map with no symbols, this is
// always bounded. However, when we have symbols, we may not be able to
// obtain a constant upper bound. Also, mapping to a negative space is
// invalid for normalization.
if (!ubConst.has_value() || *ubConst < 0) {
LLVM_DEBUG(llvm::dbgs()
<< "can't normalize map due to unknown/invalid upper bound");
// invalid for normalization. If dimension of new memrefType is dynamic,
// the value is `ShapedType::kDynamic`.
if (!ubConst.has_value())
newShape[d] = ShapedType::kDynamic;
else if (*ubConst >= 0)
newShape[d] = *ubConst + 1;
else
return memrefType;
}
// If dimension of new memrefType is dynamic, the value is -1.
newShape[d] = *ubConst + 1;
}

// Create the new memref type after trivializing the old layout map.
Expand Down
4 changes: 3 additions & 1 deletion mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -445,8 +445,10 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(func::FuncOp funcOp,
if (oldMemRefType == newMemRefType)
continue;
// TODO: Assume single layout map. Multiple maps not supported.
// TODO: Semi-affine layout not supported.
AffineMap layoutMap = oldMemRefType.getLayout().getAffineMap();
if (failed(replaceAllMemRefUsesWith(oldMemRef,
if (!layoutMap.getResult(0).isPureAffine() ||
failed(replaceAllMemRefUsesWith(oldMemRef,
/*newMemRef=*/newMemRef,
/*extraIndices=*/{},
/*indexRemap=*/layoutMap,
Expand Down
7 changes: 4 additions & 3 deletions mlir/test/Dialect/Affine/memref-bound-check.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,13 +124,14 @@ func.func @mod_floordiv_nested() {
return
}

// CHECK-LABEL: func @test_semi_affine_bailout
func.func @test_semi_affine_bailout(%N : index) {
// CHECK-LABEL: func @test_semi_affine_access
func.func @test_semi_affine_access(%N : index) {
%B = memref.alloc() : memref<10 x i32>
affine.for %i = 0 to 10 {
%idx = affine.apply affine_map<(d0)[s0] -> (d0 * s0)>(%i)[%N]
%y = affine.load %B[%idx] : memref<10 x i32>
// expected-error@-1 {{getMemRefRegion: compose affine map failed}}
// expected-error@-1 {{'affine.load' op memref out of upper bound access along dimension #1}}
// expected-error@-2 {{'affine.load' op memref out of lower bound access along dimension #1}}
}
return
}
Expand Down
Loading