[Mlir-commits] [mlir] c9e121e - [MLIR] Enhance getLargestKnownDivisor for AffineExpr floordiv/ceildiv

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2022-12-17T12:45:38+05:30
New Revision: c9e121eb44c419f0e55e6e5d7735f6c962b141bc

URL: https://github.com/llvm/llvm-project/commit/c9e121eb44c419f0e55e6e5d7735f6c962b141bc
DIFF: https://github.com/llvm/llvm-project/commit/c9e121eb44c419f0e55e6e5d7735f6c962b141bc.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/IR/AffineExpr.cpp
    mlir/test/Dialect/Affine/unroll.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 00778cd47fdb0..d783f5f6d83f5 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 c6fe6ef655fb9..582f6254facc0 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
+}