[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

mlir/lib/IR/AffineExpr.cpp

@@ -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(

mlir/test/Dialect/Affine/unroll.mlir

@@ -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() {