[MLIR] Fix AffineExpr getLargestKnownDivisor for ceildiv and floordiv

Fix AffineExpr `getLargestKnownDivisor` for ceil/floor div cases.
In these cases, nothing can be inferred on the divisor of the
result.

Add test case for `mod` as well.

Differential Revision: https://reviews.llvm.org/D112523
This commit is contained in:
Uday Bondhugula 2021-10-11 15:48:06 +05:30
parent f279e50fd0
commit 41a8b46007
2 changed files with 27 additions and 3 deletions

View File

@ -216,9 +216,11 @@ bool AffineExpr::isPureAffine() const {
int64_t AffineExpr::getLargestKnownDivisor() const {
AffineBinaryOpExpr binExpr(nullptr);
switch (getKind()) {
case AffineExprKind::SymbolId:
case AffineExprKind::CeilDiv:
LLVM_FALLTHROUGH;
case AffineExprKind::DimId:
case AffineExprKind::FloorDiv:
case AffineExprKind::SymbolId:
return 1;
case AffineExprKind::Constant:
return std::abs(this->cast<AffineConstantExpr>().getValue());
@ -229,8 +231,6 @@ int64_t AffineExpr::getLargestKnownDivisor() const {
}
case AffineExprKind::Add:
LLVM_FALLTHROUGH;
case AffineExprKind::FloorDiv:
case AffineExprKind::CeilDiv:
case AffineExprKind::Mod: {
binExpr = cast<AffineBinaryOpExpr>();
return llvm::GreatestCommonDivisor64(

View File

@ -468,6 +468,30 @@ func @loop_nest_operand2() {
return
}
// UNROLL-BY-4-LABEL: func @floordiv_mod_ub
func @floordiv_mod_ub(%M : index, %N : index) {
affine.for %i = 0 to %N step 4 {
// A cleanup should be generated here.
affine.for %j = 0 to min affine_map<(d0)[s0] -> ((16 * d0) floordiv (4 * s0))>(%i)[%N] {
"test.foo"() : () -> ()
}
}
// UNROLL-BY-4-NEXT: affine.for
// UNROLL-BY-4-NEXT: affine.for %{{.*}} = 0 to {{.*}} step 4
// UNROLL-BY-4: affine.for
affine.for %i = 0 to %N step 4 {
// No cleanup needed here.
affine.for %j = 0 to min affine_map<(d0)[s0] -> ((16 * d0) mod (4 * s0))>(%i)[%N] {
"test.foo"() : () -> ()
}
}
// UNROLL-BY-4: affine.for
// UNROLL-BY-4-NEXT: affine.for %{{.*}} = 0 to {{.*}} step 4
// UNROLL-BY-4-NOT: affine.for
// UNROLL-BY-4: return
return
}
// Difference between loop bounds is constant, but not a multiple of unroll
// factor. The cleanup loop happens to be a single iteration one and is promoted.
// UNROLL-BY-4-LABEL: func @loop_nest_operand3() {