[Mlir-commits] [mlir] [mlir] Don't assert when simplifying certain `AffineExpr`s (PR #78855)

[Felix Schneider]

https://github.com/ubfx updated https://github.com/llvm/llvm-project/pull/78855

>From 12894bbebe202a53748bbe76fb43db12be275c3e Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Sat, 20 Jan 2024 19:05:23 +0100
Subject: [PATCH 1/2] [mlir] Don't assert when simplifying certain
 `AffineExpr`s

Currently, `simplifyMul()` asserts that either `lhs` or `rhs` is symbolic or constant. This method is called by the overloaded `*` operator for `AffineExpr`s which leads to a crash when building a multiplication expression where neither operand is symbolic or constant.
This patch returns a `nullptr` from `simplifyMul()` to signal that the expression could not be simplified instead.
---
 mlir/lib/IR/AffineExpr.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 56a2d850ad3b67..94562d0f15a24a 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -774,7 +774,8 @@ static AffineExpr simplifyMul(AffineExpr lhs, AffineExpr rhs) {
     return getAffineConstantExpr(lhsConst.getValue() * rhsConst.getValue(),
                                  lhs.getContext());
 
-  assert(lhs.isSymbolicOrConstant() || rhs.isSymbolicOrConstant());
+  if (!lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant())
+    return nullptr;
 
   // Canonicalize the mul expression so that the constant/symbolic term is the
   // RHS. If both the lhs and rhs are symbolic, swap them if the lhs is a

>From b452fa8a424aef112e960bc9070790ae9b6e4889 Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Tue, 19 Mar 2024 08:08:21 +0100
Subject: [PATCH 2/2] Add Unit Test

---
 mlir/unittests/IR/AffineExprTest.cpp | 32 ++++++++++++++++++++++++++++
 mlir/unittests/IR/CMakeLists.txt     |  1 +
 2 files changed, 33 insertions(+)
 create mode 100644 mlir/unittests/IR/AffineExprTest.cpp

diff --git a/mlir/unittests/IR/AffineExprTest.cpp b/mlir/unittests/IR/AffineExprTest.cpp
new file mode 100644
index 00000000000000..d1c67e4ded750a
--- /dev/null
+++ b/mlir/unittests/IR/AffineExprTest.cpp
@@ -0,0 +1,32 @@
+//===- AffineExprTest.cpp - unit tests for affine map API -----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Builders.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+// Test creating AffineExprs using the overloaded binary operators.
+TEST(AffineExprTest, constructFromBinaryOperators) {
+  MLIRContext ctx;
+  OpBuilder b(&ctx);
+
+  auto d0 = b.getAffineDimExpr(0);
+  auto d1 = b.getAffineDimExpr(1);
+
+  auto sum = d0 + d1;
+  auto difference = d0 - d1;
+  auto product = d0 * d1;
+  auto remainder = d0 % d1;
+
+  ASSERT_EQ(sum.getKind(), AffineExprKind::Add);
+  ASSERT_EQ(difference.getKind(), AffineExprKind::Add);
+  ASSERT_EQ(product.getKind(), AffineExprKind::Mul);
+  ASSERT_EQ(remainder.getKind(), AffineExprKind::Mod);
+}
diff --git a/mlir/unittests/IR/CMakeLists.txt b/mlir/unittests/IR/CMakeLists.txt
index e7e9c3b5651693..71f8f449756ec0 100644
--- a/mlir/unittests/IR/CMakeLists.txt
+++ b/mlir/unittests/IR/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_mlir_unittest(MLIRIRTests
   AdaptorTest.cpp
+  AffineExprTest.cpp
   AffineMapTest.cpp
   AttributeTest.cpp
   DialectTest.cpp