[MLIR] Fix AffineExpr getLargestKnownDivisor for ceildiv and floordiv

[bondhugula]

Summary

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.

[ftynse]

mlir/test/Dialect/Affine/unroll.mlir

479

Nit: consider aligning check comments with the code below them.

[bondhugula]

Indent to align with code below.

mlir/lib/IR/AffineExpr.cpp

bool AffineExpr::isPureAffine() const {
   }
   llvm_unreachable("Unknown AffineExpr");
 }
 
 // Returns the greatest known integral divisor of this affine expression.
 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());
   case AffineExprKind::Mul: {
     binExpr = this->cast<AffineBinaryOpExpr>();
     return binExpr.getLHS().getLargestKnownDivisor() *
            binExpr.getRHS().getLargestKnownDivisor();
   }
   case AffineExprKind::Add:
     LLVM_FALLTHROUGH;
-  case AffineExprKind::FloorDiv:
-  case AffineExprKind::CeilDiv:
   case AffineExprKind::Mod: {
     binExpr = cast<AffineBinaryOpExpr>();
     return llvm::GreatestCommonDivisor64(
         binExpr.getLHS().getLargestKnownDivisor(),
         binExpr.getRHS().getLargestKnownDivisor());
   }
   }
   llvm_unreachable("Unknown AffineExpr");

mlir/test/Dialect/Affine/unroll.mlir

// UNROLL-BY-4-NEXT: return
   affine.for %i = 0 to 100 step 2 {
     affine.for %j = affine_map<(d0) -> (d0)> (%i) to affine_map<(d0) -> (5*d0 + 4)> (%i) {
       %x = "foo"() : () -> i32
     }
   }
   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() {
 func @loop_nest_operand3() {
   // UNROLL-BY-4: affine.for %arg0 = 0 to 100 step 2 {
   affine.for %i = 0 to 100 step 2 {
     // UNROLL-BY-4: affine.for %arg1 = [[$MAP11]](%arg0) to #map{{[0-9]+}}(%arg0) step 4 {
     // UNROLL-BY-4-NEXT: %1 = "foo"() : () -> i32