[MLIR] Simplify semi-affine expressions

[yash]

Summary

Simplify semi-affine expression for the operations like ceildiv, floordiv and modulo by any given symbol by checking divisibilty by that symbol.

Some properties used in simplification are:

1. Commutative property of the floordiv and ceildiv:

((expr1 floordiv expr2) floordiv expr3 ) = ((expr1 floordiv expr3) floordiv expr2)
((expr1 ceildiv expr2) ceildiv expr3 ) = ((expr1 ceildiv expr3) ceildiv expr2)
While simplification if operations are different no simplification is possible as there is no property that simplify expressions like these: ((expr1 ceildiv expr2) floordiv expr3) or ((expr1 floordiv expr2) ceildiv expr3).

2. If both expr1 and expr2 are divisible by the expr3 then:

(expr1 % expr2) / expr3 = ((expr1 / expr3) % (expr2 / expr3)) "where / is divide symbol".

3. If expr1 is divisible by expr2 then expr1 % expr2 = 0.

Signed-off-by: Yash Jain <[e-mail]>

[yash]

Minor update.

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

255–256

This is a bit fuzzy. If divOpKind can be mod as well, then we need a different name - may be divModOpKind. If it can be only floordiv, ceildiv, and mod, you need to also assert the precondition as the first thing.

258

'0' is divisible by all symbols (since all divisor symbols are expected to be positive and thus non-zero). So, this should be returning true for zero constant?

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

273–275

This looks incorrect - shouldn't you be passing divOpKind as the third argument?

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

305

Nit: nullptr -> null expression

386–390

The third argument for isDivisibleSymbol is always floor, ceil, or mod. Should assert at the method?

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

249

Nit: signifies here ... -> specifies what kind of division or mod operation associates the two.

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

287

Nit: with -> in

288

No need of $ prefix for map?

323

Terminate with a ful stop. simplify -> simplifies

[yash]

divOpKind changed to opKind.
Constant 0 can be divided handled.

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

264–266

Doesn't look there is a test case to cover this - please add one.

326

Nit: prefer to make it explicit: ... != 0

[bondhugula]

But no property for the expressions with different operations ((expr1 ceildiv expr2) floordiv expr3)

Please rephrase this.

[yash]

Added test case for division of constant 0 by symbol.

[bondhugula]

In D84920#2186529, @yash wrote:

Added test case for division of constant 0 by symbol.

I can't find the test case. Did you update the patch?

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

1029–1031

-> // Simplify semi-affine expressions separately.

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

330

Drop extra line at end of file. Please use git diff --check HEAD~ to check for whitespace errors.

[yash]

Removed the extra line and edited some comments.

[bondhugula]

Thanks for revising. Looks good to me - minor comments.

mlir/lib/IR/AffineExpr.cpp

249–250

Sentence is broken.

252

expr -> `expr`

258

Nit: THe -> The

315

You can drop the "useful for ...". This is just implementing the said functionality.

316–317

Nit: Rephrase A null expression is returned ...

[yash]

Minor update in comments.

mlir/lib/IR/AffineExpr.cpp

AffineExpr AffineBinaryOpExpr::getRHS() const {
   return static_cast<ImplType *>(expr)->rhs;
 }
 
 AffineDimExpr::AffineDimExpr(AffineExpr::ImplType *ptr) : AffineExpr(ptr) {}
 unsigned AffineDimExpr::getPosition() const {
   return static_cast<ImplType *>(expr)->position;
 }
 
+/// Returns true if the expression is divisible by the given symbol with
+/// position `symbolPos`. The argument `divOpKind` signifies here which
+/// operation incorporates this division. It helps in implementing the
+/// commutative property of the floordiv and ceildiv operations. If the argument
+/// `exprKind` is floordiv and expr is also a binary expression of a floordiv
+/// operation, then the commutative property can be used otherwise, the floordiv
+/// operation is not divisible. The same argument holds for ceildiv operation.
+static bool isDivisibleBySymbol(AffineExpr expr, unsigned symbolPos,
+                                AffineExprKind divOpKind) {
+  switch (expr.getKind()) {
+  case AffineExprKind::Constant:
+  case AffineExprKind::DimId:
+    return false;
+  case AffineExprKind::SymbolId:
+    return (expr.cast<AffineSymbolExpr>().getPosition() == symbolPos);
+  // Checks divisibility by the given symbol for both operands.
+  case AffineExprKind::Add: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, divOpKind) &&
+           isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, divOpKind);
+  }
+  // Checks divisibility by the given symbol for both operands.
+  case AffineExprKind::Mod: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos,
+                               AffineExprKind::Mod) &&
+           isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos,
+                               AffineExprKind::Mod);
+  }
+  // Checks if any of the operand divisible by the given symbol.
+  case AffineExprKind::Mul: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, divOpKind) ||
+           isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, divOpKind);
+  }
+  // Floordiv and ceildiv are divisible by the given symbol when the first
+  // operand is divisible, and the affine expression kind of the argument expr
+  // is same as the argument `divOpKind`. This can be inferred from commutative
+  // property of floordiv and ceildiv operations and are as follow:
+  // (exp1 floordiv exp2) floordiv exp3 = (exp1 floordiv exp3) floordiv exp2
+  // (exp1 ceildiv exp2) ceildiv exp3 = (exp1 ceildiv exp3) ceildiv expr2
+  // It will fail if operations are not same. For example:
+  // (exps1 ceildiv exp2) floordiv exp3 can not be simplified.
+  case AffineExprKind::FloorDiv:
+  case AffineExprKind::CeilDiv: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    if (divOpKind != expr.getKind())
+      return false;
+    return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind());
+  }
+  }
+  llvm_unreachable("Unknown AffineExpr");
+}
+
+/// Divides the given expression by the given symbol at position `symbolPos`,
+/// useful for simplifying the semi-affine expressions. It considers
+/// divisibility condition is checked before calling itself. It returns a
+/// nullptr whenever the divisibility condition fails.
+static AffineExpr symbolicDivide(AffineExpr expr, unsigned symbolPos,
+                                 AffineExprKind divOpKind) {
+  switch (expr.getKind()) {
+  case AffineExprKind::Constant:
+  case AffineExprKind::DimId:
+    return nullptr;
+  case AffineExprKind::SymbolId:
+    return getAffineConstantExpr(1, expr.getContext());
+  // Dividing both operands by the given symbol.
+  case AffineExprKind::Add: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    return getAffineBinaryOpExpr(
+        expr.getKind(),
+        symbolicDivide(binaryExpr.getLHS(), symbolPos, divOpKind),
+        symbolicDivide(binaryExpr.getRHS(), symbolPos, divOpKind));
+  }
+  // Dividing both operands by the given symbol.
+  case AffineExprKind::Mod: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    return getAffineBinaryOpExpr(
+        expr.getKind(),
+        symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()),
+        symbolicDivide(binaryExpr.getRHS(), symbolPos, expr.getKind()));
+  }
+  // Dividing any of the operand by the given symbol.
+  case AffineExprKind::Mul: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, divOpKind))
+      return binaryExpr.getLHS() *
+             symbolicDivide(binaryExpr.getRHS(), symbolPos, divOpKind);
+    return symbolicDivide(binaryExpr.getLHS(), symbolPos, divOpKind) *
+           binaryExpr.getRHS();
+  }
+  // Dividing first operand only by the given symbol.
+  case AffineExprKind::FloorDiv:
+  case AffineExprKind::CeilDiv: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    return getAffineBinaryOpExpr(
+        expr.getKind(),
+        symbolicDivide(binaryExpr.getLHS(), symbolPos, expr.getKind()),
+        binaryExpr.getRHS());
+  }
+  }
+  llvm_unreachable("Unknown AffineExpr");
+}
+
+/// Simplify a semi-affine expression by handling modulo, floordiv, or ceildiv
+/// operations when the second operand simplifies to a symbol and the first
+/// operand is divisible by that symbol. It can be applied to any semi-affine
+/// expression. Returned expression can either be a semi-affine or pure affine
+/// expression.
+static AffineExpr simplifySemiAffine(AffineExpr expr) {
+  switch (expr.getKind()) {
+  case AffineExprKind::Constant:
+  case AffineExprKind::DimId:
+  case AffineExprKind::SymbolId:
+    return expr;
+  case AffineExprKind::Add:
+  case AffineExprKind::Mul: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    return getAffineBinaryOpExpr(expr.getKind(),
+                                 simplifySemiAffine(binaryExpr.getLHS()),
+                                 simplifySemiAffine(binaryExpr.getRHS()));
+  }
+  // Check if the simplification of the second operand is a symbol, and the
+  // first operand is divisible by it. If the operation is a modulo, a constant
+  // zero expression is returned. In the case of floordiv and ceildiv, the
+  // symbol from the simplification of the second operand divides the first
+  // operand. Otherwise, simplification is not possible.
+  case AffineExprKind::FloorDiv:
+  case AffineExprKind::CeilDiv:
+  case AffineExprKind::Mod: {
+    AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+    AffineExpr sLHS = simplifySemiAffine(binaryExpr.getLHS());
+    AffineExpr sRHS = simplifySemiAffine(binaryExpr.getRHS());
+    AffineSymbolExpr symbolExpr =
+        simplifySemiAffine(binaryExpr.getRHS()).dyn_cast<AffineSymbolExpr>();
+    if (!symbolExpr)
+      return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
+    unsigned symbolPos = symbolExpr.getPosition();
+    if (!isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind()))
+      return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
+    if (expr.getKind() == AffineExprKind::Mod)
+      return getAffineConstantExpr(0, expr.getContext());
+    return symbolicDivide(sLHS, symbolPos, expr.getKind());
+  }
+  }
+  llvm_unreachable("Unknown AffineExpr");
+}
+
 static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position,
                                        MLIRContext *context) {
   auto assignCtx = [context](AffineDimExprStorage *storage) {
     storage->context = context;
   };
 
   StorageUniquer &uniquer = context->getAffineUniquer();
   return uniquer.get<AffineDimExprStorage>(
int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) {
   if ((it = llvm::find(localExprs, localExpr)) == localExprs.end())
     return -1;
   return it - localExprs.begin();
 }
 
 /// Simplify the affine expression by flattening it and reconstructing it.
 AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims,
                                     unsigned numSymbols) {
-  // TODO: only pure affine for now. The simplification here can
-  // be extended to semi-affine maps in the future.
+  // Simplify the semi-affine expression.
+  if (!expr.isPureAffine())
+    expr = simplifySemiAffine(expr);
   if (!expr.isPureAffine())
     return expr;
 
   SimpleAffineExprFlattener flattener(numDims, numSymbols);
   flattener.walkPostOrder(expr);
   ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
   auto simplifiedExpr =
       getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols,
                                 flattener.localExprs, expr.getContext());
   flattener.operandExprStack.pop_back();
   assert(flattener.operandExprStack.empty());
 
   return simplifiedExpr;
 }

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

 
 // CHECK-DAG: #[[MAP_0D:.*]] = affine_map<() -> ()>
 
 // CHECK-LABEL: func @simplify_zero_dim_map
 func @simplify_zero_dim_map(%in : memref<f32>) -> f32 {
   %out = affine.load %in[] : memref<f32>
   return %out : f32
 }
+
+// -----
+
+// Tests simplification of a semi-affine expression with various cases.
+// CHECK-DAG: #[[$map0:.*]] = affine_map<()[s0, s1] -> (-(s1 floordiv s0) + 2)>
+// CHECK-DAG: #[[$map1:.*]] = affine_map<()[s0, s1] -> (-(s1 floordiv s0) + 42)>
+
+// Tests simplification of a semi-affine expression with a modulo operartion on a floordiv and multiplication.
+// CHECK-LABEL: func @semiaffine_mod
+func @semiaffine_mod(%arg0: index, %arg1: index) -> index {
+  %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + s0 * s0) mod  s0)> (%arg0)[%arg1]
+  // CHECK:       %[[CST:.*]] = constant 0
+  return %a : index
+}
+
+// Tests simplification of a semi-affine expression with a nested floordiv and a floordiv on modulo operation.
+// CHECK-LABEL: func @semiaffine_floordiv
+func @semiaffine_floordiv(%arg0: index, %arg1: index) -> index {
+  %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + ((2 * s0) mod (3 * s0))) floordiv s0)> (%arg0)[%arg1]
+  // CHECK: affine.apply #[[$map0]]()[%arg1, %arg0]
+  return %a : index
+}
+
+// Tests simplification of a semi-affine expression with a ceildiv operation.
+// CHECK-LABEL: func @semiaffine_ceildiv
+func @semiaffine_ceildiv(%arg0: index, %arg1: index) -> index {
+  %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + s0 * 42) ceildiv  s0)> (%arg0)[%arg1]
+  // CHECK: affine.apply #[[$map1]]()[%arg1, %arg0]
+  return %a : index
+}
+
+// Tests simplification of a semi-affine expression with a nested ceildiv operation and further simplifications after performing ceildiv.
+// CHECK-LABEL: func @semiaffine_composite_floor
+func @semiaffine_composite_floor(%arg0: index, %arg1: index) -> index {
+  %a = affine.apply affine_map<(d0)[s0] ->(((((s0 * 2) ceildiv 4) * 5) + s0 * 42) ceildiv  s0)> (%arg0)[%arg1]
+  // CHECK:       %[[CST:.*]] = constant 47
+  return %a : index
+}
+
+// Tests simplification of a semi-affine expression with a modulo operation with a second operand that simplify to symbol.
+// CHECK-LABEL: func @semiaffine_unsimplified_symbol
+func @semiaffine_unsimplified_symbol(%arg0: index, %arg1: index) -> index {
+  %a = affine.apply affine_map<(d0)[s0] ->(s0 mod (2 * s0 - s0))> (%arg0)[%arg1]
+  // CHECK:       %[[CST:.*]] = constant 0
+  return %a : index
+}
+