[mlir][affine] remove divide zero check when simplifer affineMap (#64622)

[lipracer]

When performing constant folding on the affineApplyOp, there is a division of 0 in the affine map.
related issue

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir-affine

Changes

When performing constant folding on the affineApplyOp, there is a division of 0 in the affine map

---

Full diff: https://github.com/llvm/llvm-project/pull/68519.diff

4 Files Affected:

• (modified) mlir/include/mlir/IR/AffineExprVisitor.h (+2)
• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+7)
• (modified) mlir/lib/IR/AffineExpr.cpp (-1)
• (modified) mlir/lib/IR/AffineMap.cpp (+28-5)

diff --git a/mlir/include/mlir/IR/AffineExprVisitor.h b/mlir/include/mlir/IR/AffineExprVisitor.h
index f6216614c2238e1..aa2484e8fe182aa 100644
--- a/mlir/include/mlir/IR/AffineExprVisitor.h
+++ b/mlir/include/mlir/IR/AffineExprVisitor.h
@@ -289,6 +289,8 @@ class SimpleAffineExprFlattener
   // A mod expression "expr mod c" is thus flattened by introducing a new local
   // variable q (= expr floordiv c), such that expr mod c is replaced with
   // 'expr - c * q' and c * q <= expr <= c * q + c - 1 are added to localVarCst.
+  // TODO Modify the return value to LogicResult and handle cases where the
+  // division is zero
   void visitModExpr(AffineBinaryOpExpr expr);
 
 protected:
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 113f4cfc31c104b..292795673cfc659 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -569,6 +569,13 @@ bool AffineApplyOp::isValidSymbol(Region *region) {
 }
 
 OpFoldResult AffineApplyOp::fold(FoldAdaptor adaptor) {
+  ScopedDiagnosticHandler foldDiagnosticHandler(
+      getContext(), [this](Diagnostic &diag) {
+        diag.attachNote(getLoc())
+            .append("see current operation: ")
+            .appendOp(*getOperation(), OpPrintingFlags().printGenericOpForm());
+        return failure();
+      });
   auto map = getAffineMap();
 
   // Fold dims and symbols to existing values.
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 7eccbca4e6e7a1a..1340457833788d0 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -511,7 +511,6 @@ unsigned AffineSymbolExpr::getPosition() const {
 
 AffineExpr mlir::getAffineSymbolExpr(unsigned position, MLIRContext *context) {
   return getAffineDimOrSymbol(AffineExprKind::SymbolId, position, context);
-  ;
 }
 
 AffineConstantExpr::AffineConstantExpr(AffineExpr::ImplType *ptr)
diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 9cdac964710ca86..ff9cd5d6ff1fc56 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -11,6 +11,7 @@
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/STLExtras.h"
@@ -56,13 +57,34 @@ class AffineExprConstantFolder {
           expr, [](int64_t lhs, int64_t rhs) { return lhs * rhs; });
     case AffineExprKind::Mod:
       return constantFoldBinExpr(
-          expr, [](int64_t lhs, int64_t rhs) { return mod(lhs, rhs); });
+          expr, [expr](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {
+            if (rhs < 1) {
+              emitWarning(UnknownLoc::get(expr.getContext()),
+                          "mod division by zero!");
+              return std::nullopt;
+            }
+            return mod(lhs, rhs);
+          });
     case AffineExprKind::FloorDiv:
       return constantFoldBinExpr(
-          expr, [](int64_t lhs, int64_t rhs) { return floorDiv(lhs, rhs); });
+          expr, [expr](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {
+            if (0 == rhs) {
+              emitWarning(UnknownLoc::get(expr.getContext()),
+                          "floor division by zero!");
+              return std::nullopt;
+            }
+            return floorDiv(lhs, rhs);
+          });
     case AffineExprKind::CeilDiv:
       return constantFoldBinExpr(
-          expr, [](int64_t lhs, int64_t rhs) { return ceilDiv(lhs, rhs); });
+          expr, [expr](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {
+            if (0 == rhs) {
+              emitWarning(UnknownLoc::get(expr.getContext()),
+                          "ceil division by zero!");
+              return std::nullopt;
+            }
+            return ceilDiv(lhs, rhs);
+          });
     case AffineExprKind::Constant:
       return expr.cast<AffineConstantExpr>().getValue();
     case AffineExprKind::DimId:
@@ -81,8 +103,9 @@ class AffineExprConstantFolder {
   }
 
   // TODO: Change these to operate on APInts too.
-  std::optional<int64_t> constantFoldBinExpr(AffineExpr expr,
-                                             int64_t (*op)(int64_t, int64_t)) {
+  std::optional<int64_t> constantFoldBinExpr(
+      AffineExpr expr,
+      llvm::function_ref<std::optional<int64_t>(int64_t, int64_t)> op) {
     auto binOpExpr = expr.cast<AffineBinaryOpExpr>();
     if (auto lhs = constantFoldImpl(binOpExpr.getLHS()))
       if (auto rhs = constantFoldImpl(binOpExpr.getRHS()))

[matthias-springer]

I have not looked at the fix in detail yet, just one comment: This bug looks like a case of "folding leads to invalid op". But we do not create an invalid op but already crash when trying to create it. Similar issues were discussed in the context of "poison semantics". You could have IR such as:

scf.if %c {
  // The folder of affine.apply crashes due to a division by 0, but %c is such that the branch would never be executed if there would be a division by zero. E.g., %c could be the result of a size check.
  affine.apply
}

I think the consensus on such issues was, if a folder that turns dynamic information into static information (I think that's what this folder here does) would create an invalid op, then the op should not be folded. Same in this case, we should not fold the op, but also not crash. We probably also don't want to emit warnings.

[lipracer]

Thanks,I am found the poison-sem https://mlir.llvm.org/OpenMeetings/2022-11-03-Poison-Semantics-For-MLIR.pdf, just remove runtime check and replace foldedop with freezeop.
Affine expression simplifier eg. a mod b <==> a - N*b,maybe we need to define a poison value to affine express.affine express itself always well-def,and simplifier also well but the poison value is part of express.
c compiler like gcc always report a warning when division is 0.

[joker-eph]

Can you please update the PR title to be descriptive? See https://mlir.llvm.org/getting_started/Contributing/#commit-messages

[lipracer]

When AffineApplyOp::fold return PoisonAttr, the test case crash and tell me UBDialect need to initilize.

[matthias-springer]

When AffineApplyOp::fold return PoisonAttr, the test case crash and tell me UBDialect need to initilize.

That means the UB dialect would have to be declared as a dependent dialect of the Affine dialect (because such ops are created by the canonicalization). I am not an expert on poison semantics, I'd ask on Discord to see if that's desirable and the recommended path. Also, @kuhar may have an opinion on this.

[lipracer]

Perhaps the arith dialect can fold out 'poison' semantics, is this pull request ready to be merged?

[matthias-springer]

I'm a bit hesitant to add a LogicalResult return value to the entire visitor logic here. It is not clear to me whether this operation should fold to a poison value in case of a division by zero. If that's not the case, we won't need the LogicalResult return value. Until that's clear, I would just go with a simpler fix: just stop folding when there's a division by zero, by treating it the same was as "not a constant".

[matthias-springer]

That being said, don't throw this code away yet, if the consensus is that folding/canonicalization should create poison values, we probably need exactly the code that you wrote here.

[lipracer]

@matthias-springer I have found a fix for not being able to create ‘poisonstr’. ‘ubdialect’ is lazy to load and seems to only occur during the parse phase, so I can load it during the initialization phase.

[lipracer]

@joker-eph Can you help me review it?

[jsetoain]

Looks good to me. Wait for a couple more days in case the other reviewers have anything else to say, please. It is not a trivial change, and I may have overlooked something.