[Mlir-commits] [mlir] [mlir] fix a crash (PR #68519)

[llvmlistbot at llvm.org]

https://github.com/lipracer created https://github.com/llvm/llvm-project/pull/68519

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

>From a1c07d5aa9b3828d2fa172f1279f61303577efa9 Mon Sep 17 00:00:00 2001
From: lipracer <lipracer at gmail.com>
Date: Sun, 8 Oct 2023 18:31:59 +0800
Subject: [PATCH] [mlir] fix a crash

When performing constant folding on the affineApplyOp,
there is a division of 0 in the affine map
---
 mlir/include/mlir/IR/AffineExprVisitor.h |  2 ++
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp |  7 +++++
 mlir/lib/IR/AffineExpr.cpp               |  1 -
 mlir/lib/IR/AffineMap.cpp                | 33 ++++++++++++++++++++----
 4 files changed, 37 insertions(+), 6 deletions(-)

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()))