mirror of
https://github.com/intel/llvm.git
synced 2026-01-26 12:26:52 +08:00
[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:
@@ -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: {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user