[Mlir-commits] [mlir] [mlir][affine] Use affine symbol for RHS of `divMod` expressions (PR #107976)

[Felix Schneider]

https://github.com/ubfx created https://github.com/llvm/llvm-project/pull/107976

Formally, (semi) `AffineExpr`s require the RHS of `(ceil|floor)div` and `mod` binary operations to be symbolic or constant. This is currently not enforced, which can make analyses of `AffineExpr`s more complicated, e.g. in the `ValueBoundsOpInterface` (see https://github.com/llvm/llvm-project/pull/103010#issuecomment-2285683858).

This patch changes the affine utility function `getDivMod()` to create its `AffineExpr`s with a symbolic RHS such that no illegal IR is created there.

>From 9e94b95195b193b318d261c490362addc331b28f Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Tue, 10 Sep 2024 08:33:04 +0200
Subject: [PATCH] [mlir][affine] Use affine symbol for RHS of `divMod`
 expressions

Formally, (semi) `AffineExpr`s require the RHS of `(ceil|floor)div` and
`mod` binary operations to be symbolic or constant. This is currently not
enforced, which can make analyses of `AffineExpr`s more complicated, e.g.
in the `ValueBoundsOpInterface` (see https://github.com/llvm/llvm-project/pull/103010#issuecomment-2285683858).

This patch changes the affine utility function `getDivMod()` to create
its `AffineExpr`s with a symbolic RHS such that no illegal IR is created
there.
---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp               | 8 ++++----
 mlir/test/Dialect/Affine/affine-expand-index-ops.mlir | 8 ++++----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 898467d573362b..5b84d5c1c1d8fa 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1826,12 +1826,12 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
 DivModValue mlir::affine::getDivMod(OpBuilder &b, Location loc, Value lhs,
                                     Value rhs) {
   DivModValue result;
-  AffineExpr d0, d1;
-  bindDims(b.getContext(), d0, d1);
+  AffineExpr d0 = b.getAffineDimExpr(0);
+  AffineExpr s0 = b.getAffineSymbolExpr(0);
   result.quotient =
-      affine::makeComposedAffineApply(b, loc, d0.floorDiv(d1), {lhs, rhs});
+      affine::makeComposedAffineApply(b, loc, d0.floorDiv(s0), {lhs, rhs});
   result.remainder =
-      affine::makeComposedAffineApply(b, loc, d0 % d1, {lhs, rhs});
+      affine::makeComposedAffineApply(b, loc, d0 % s0, {lhs, rhs});
   return result;
 }
 
diff --git a/mlir/test/Dialect/Affine/affine-expand-index-ops.mlir b/mlir/test/Dialect/Affine/affine-expand-index-ops.mlir
index 70b7f397ad4fec..668ab8056f2004 100644
--- a/mlir/test/Dialect/Affine/affine-expand-index-ops.mlir
+++ b/mlir/test/Dialect/Affine/affine-expand-index-ops.mlir
@@ -21,8 +21,8 @@ func.func @static_basis(%linear_index: index) -> (index, index, index) {
 // -----
 
 //   CHECK-DAG:   #[[$map0:.+]] = affine_map<()[s0, s1, s2] -> (s2 floordiv (s0 * s1))>
-//   CHECK-DAG:   #[[$map1:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s0 * s1)) floordiv s1)>
-//   CHECK-DAG:   #[[$map2:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s0 * s1)) mod s1)>
+//   CHECK-DAG:   #[[$map1:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s1 * s0)) floordiv s0)>
+//   CHECK-DAG:   #[[$map2:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s1 * s0)) mod s0)>
 
 // CHECK-LABEL: @dynamic_basis
 //  CHECK-SAME:    (%[[IDX:.+]]: index, %[[MEMREF:.+]]: memref
@@ -31,8 +31,8 @@ func.func @static_basis(%linear_index: index) -> (index, index, index) {
 //        CHECK:  %[[DIM1:.+]] = memref.dim %[[MEMREF]], %[[C1]] :
 //        CHECK:  %[[DIM2:.+]] = memref.dim %[[MEMREF]], %[[C2]] :
 //       CHECK:   %[[N:.+]] = affine.apply #[[$map0]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]
-//       CHECK:   %[[P:.+]] = affine.apply #[[$map1]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]
-//       CHECK:   %[[Q:.+]] = affine.apply #[[$map2]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]
+//       CHECK:   %[[P:.+]] = affine.apply #[[$map1]]()[%[[DIM2]], %[[DIM1]], %[[IDX]]]
+//       CHECK:   %[[Q:.+]] = affine.apply #[[$map2]]()[%[[DIM2]], %[[DIM1]], %[[IDX]]]
 //       CHECK:   return %[[N]], %[[P]], %[[Q]]
 func.func @dynamic_basis(%linear_index: index, %src: memref<?x?x?xf32>) -> (index, index, index) {
   %c0 = arith.constant 0 : index