[MLIR] Enhance getLargestKnownDivisor for AffineExpr floordiv/ceildiv

The largest known divisor for expressions like (32 * d0 + 32, 128)
ceildiv 8 wasn't being computed tightly; a conservative value of 1 was
being returned. Address this. This leads to a broad improvement for
several affine analyses and rewrites that depend on knowing whether
something is a multiple of a specific constant or such largest known
constant.

Differential Revision: https://reviews.llvm.org/D140185
This commit is contained in:
Uday Bondhugula
2022-12-15 20:54:19 +05:30
parent 85956bd552
commit c9e121eb44
2 changed files with 38 additions and 4 deletions

View File

@@ -219,12 +219,25 @@ bool AffineExpr::isPureAffine() const {
int64_t AffineExpr::getLargestKnownDivisor() const {
AffineBinaryOpExpr binExpr(nullptr);
switch (getKind()) {
case AffineExprKind::CeilDiv:
[[fallthrough]];
case AffineExprKind::DimId:
case AffineExprKind::FloorDiv:
[[fallthrough]];
case AffineExprKind::SymbolId:
return 1;
case AffineExprKind::CeilDiv:
[[fallthrough]];
case AffineExprKind::FloorDiv: {
// If the RHS is a constant and divides the known divisor on the LHS, the
// quotient is a known divisor of the expression.
binExpr = this->cast<AffineBinaryOpExpr>();
auto rhs = binExpr.getRHS().dyn_cast<AffineConstantExpr>();
// Leave alone undefined expressions.
if (rhs && rhs.getValue() != 0) {
int64_t lhsDiv = binExpr.getLHS().getLargestKnownDivisor();
if (lhsDiv % rhs.getValue() == 0)
return lhsDiv / rhs.getValue();
}
return 1;
}
case AffineExprKind::Constant:
return std::abs(this->cast<AffineConstantExpr>().getValue());
case AffineExprKind::Mul: {

View File

@@ -746,4 +746,25 @@ func.func @unroll_cleanup_loop_with_identical_unroll_factor() {
// UNROLL-CLEANUP-LOOP-NEXT: %[[V4:.*]] = affine.apply {{.*}}
// UNROLL-CLEANUP-LOOP-NEXT: {{.*}} = "foo"(%[[V4]]) : (index) -> i32
// UNROLL-CLEANUP-LOOP-NEXT: return
}
}
// UNROLL-BY-4-LABEL: func @known_multiple_ceildiv
func.func @known_multiple_ceildiv(%N: index, %S: index) {
%cst = arith.constant 0.0 : f32
%m = memref.alloc(%S) : memref<?xf32>
// This exercises affine expr getLargestKnownDivisor for the ceildiv case.
affine.for %i = 0 to affine_map<(d0) -> (32 * d0 + 64)>(%N) step 8 {
affine.store %cst, %m[%i] : memref<?xf32>
}
// UNROLL-BY-4: affine.for %{{.*}} = 0 to {{.*}} step 32
// UNROLL-BY-4-NOT: affine.for
// This exercises affine expr getLargestKnownDivisor for floordiv.
affine.for %i = 0 to affine_map<(d0) -> ((32 * d0 + 64) floordiv 8)>(%N) {
affine.store %cst, %m[%i] : memref<?xf32>
}
// UNROLL-BY-4: affine.for %{{.*}} = 0 to {{.*}} step 4
// UNROLL-BY-4-NOT: affine.for
// UNROLL-BY-4: return
return
}