[Mlir-commits] [mlir] 41a8b46 - [MLIR] Fix AffineExpr getLargestKnownDivisor for ceildiv and floordiv

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2021-10-26T16:21:29+05:30
New Revision: 41a8b4600731b0ce702adb6087c73133089c78d5

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

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

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 2e13d6ae62c4..b98731031521 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/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(

diff  --git a/mlir/test/Dialect/Affine/unroll.mlir b/mlir/test/Dialect/Affine/unroll.mlir
index 4b302c45b036..fe7b8e4c10ad 100644
--- a/mlir/test/Dialect/Affine/unroll.mlir
+++ b/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() {