diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp --- a/mlir/lib/IR/AffineExpr.cpp +++ b/mlir/lib/IR/AffineExpr.cpp @@ -216,9 +216,11 @@ 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().getValue()); @@ -229,8 +231,6 @@ } case AffineExprKind::Add: LLVM_FALLTHROUGH; - case AffineExprKind::FloorDiv: - case AffineExprKind::CeilDiv: case AffineExprKind::Mod: { binExpr = cast(); return llvm::GreatestCommonDivisor64( diff --git a/mlir/test/Dialect/Affine/unroll.mlir b/mlir/test/Dialect/Affine/unroll.mlir --- a/mlir/test/Dialect/Affine/unroll.mlir +++ b/mlir/test/Dialect/Affine/unroll.mlir @@ -468,6 +468,30 @@ 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() {