[MLIR] Enhance getLargestKnownDivisor for AffineExpr floordiv/ceildiv

[bondhugula]

Summary

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.

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::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: {
     binExpr = this->cast<AffineBinaryOpExpr>();
     return binExpr.getLHS().getLargestKnownDivisor() *
            binExpr.getRHS().getLargestKnownDivisor();
   }
   case AffineExprKind::Add:

mlir/test/Dialect/Affine/unroll.mlir

 // UNROLL-CLEANUP-LOOP-NEXT: {{.*}} = "foo"(%[[V1]]) : (index) -> i32
 // UNROLL-CLEANUP-LOOP-NEXT: %[[V2:.*]] = affine.apply {{.*}}
 // UNROLL-CLEANUP-LOOP-NEXT: {{.*}} = "foo"(%[[V2]]) : (index) -> i32
 // UNROLL-CLEANUP-LOOP-NEXT: %[[V3:.*]] = affine.apply {{.*}}
 // UNROLL-CLEANUP-LOOP-NEXT: {{.*}} = "foo"(%[[V3]]) : (index) -> i32
 // 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
+}