From c9e121eb44c419f0e55e6e5d7735f6c962b141bc Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday@polymagelabs.com>
Date: Thu, 15 Dec 2022 20:54:19 +0530
Subject: [PATCH] [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
---
 mlir/lib/IR/AffineExpr.cpp           | 19 ++++++++++++++++---
 mlir/test/Dialect/Affine/unroll.mlir | 23 ++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 00778cd..d783f5f 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -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: {
diff --git a/mlir/test/Dialect/Affine/unroll.mlir b/mlir/test/Dialect/Affine/unroll.mlir
index c6fe6ef..582f625 100644
--- a/mlir/test/Dialect/Affine/unroll.mlir
+++ b/mlir/test/Dialect/Affine/unroll.mlir
@@ -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
-}
\ No newline at end of file
+}
+
+// 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
+}
-- 
2.7.4