Fix bug in visitDivExpr, visitMulExpr and visitModExpr

Whenever the result of a div or mod affine expression is a
constant expression, place the value in the constant index of the
flattened expression instead of adding it as a local expression.

Author: arnab-polymage

mlir/lib/IR/AffineExpr.cpp

@@ -1177,10 +1177,9 @@ static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
     if (flatExprs[numDims + numSymbols + it.index()] == 0)
       continue;
     AffineExpr expr = it.value();
-    auto binaryExpr = dyn_cast<AffineBinaryOpExpr>(expr);
-    if (!binaryExpr)
-      continue;
-
+    // Local expression cannot be a dimension, symbol or a constant -- it
+    // should be a binary op expression.
+    auto binaryExpr = cast<AffineBinaryOpExpr>(expr);
     AffineExpr lhs = binaryExpr.getLHS();
     AffineExpr rhs = binaryExpr.getRHS();
     if (!((isa<AffineDimExpr>(lhs) || isa<AffineSymbolExpr>(lhs)) &&
@@ -1295,7 +1294,23 @@ LogicalResult SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) {
                                              localExprs, context);
     AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
                                              localExprs, context);
-    return addLocalVariableSemiAffine(mulLhs, rhs, a * b, lhs, lhs.size());
+    AffineExpr mulExpr = a * b;
+    if (auto constMulExpr = dyn_cast<AffineConstantExpr>(mulExpr)) {
+      std::fill(lhs.begin(), lhs.end(), 0);
+      lhs[getConstantIndex()] = constMulExpr.getValue();
+      return success();
+    }
+    if (auto dimMulExpr = dyn_cast<AffineDimExpr>(mulExpr)) {
+      std::fill(lhs.begin(), lhs.end(), 0);
+      lhs[getDimStartIndex() + dimMulExpr.getPosition()] = 1;
+      return success();
+    }
+    if (auto symbolMulExpr = dyn_cast<AffineSymbolExpr>(mulExpr)) {
+      std::fill(lhs.begin(), lhs.end(), 0);
+      lhs[getSymbolStartIndex() + symbolMulExpr.getPosition()] = 1;
+      return success();
+    }
+    return addLocalVariableSemiAffine(mulLhs, rhs, mulExpr, lhs, lhs.size());
   }
 
   // Get the RHS constant.
@@ -1348,6 +1363,21 @@ LogicalResult SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) {
     AffineExpr divisorExpr = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
                                                        localExprs, context);
     AffineExpr modExpr = dividendExpr % divisorExpr;
+    if (auto constModExpr = dyn_cast<AffineConstantExpr>(modExpr)) {
+      std::fill(lhs.begin(), lhs.end(), 0);
+      lhs[getConstantIndex()] = constModExpr.getValue();
+      return success();
+    }
+    if (auto dimModExpr = dyn_cast<AffineDimExpr>(modExpr)) {
+      std::fill(lhs.begin(), lhs.end(), 0);
+      lhs[getDimStartIndex() + dimModExpr.getPosition()] = 1;
+      return success();
+    }
+    if (auto symbolModExpr = dyn_cast<AffineSymbolExpr>(modExpr)) {
+      std::fill(lhs.begin(), lhs.end(), 0);
+      lhs[getSymbolStartIndex() + symbolModExpr.getPosition()] = 1;
+      return success();
+    }
     return addLocalVariableSemiAffine(modLhs, rhs, modExpr, lhs, lhs.size());
   }
 
@@ -1482,6 +1512,21 @@ LogicalResult SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr,
     AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
                                              localExprs, context);
     AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
+    if (auto constDivExpr = dyn_cast<AffineConstantExpr>(divExpr)) {
+      std::fill(lhs.begin(), lhs.end(), 0);
+      lhs[getConstantIndex()] = constDivExpr.getValue();
+      return success();
+    }
+    if (auto dimDivExpr = dyn_cast<AffineDimExpr>(divExpr)) {
+      std::fill(lhs.begin(), lhs.end(), 0);
+      lhs[getDimStartIndex() + dimDivExpr.getPosition()] = 1;
+      return success();
+    }
+    if (auto symbolDivExpr = dyn_cast<AffineSymbolExpr>(divExpr)) {
+      std::fill(lhs.begin(), lhs.end(), 0);
+      lhs[getSymbolStartIndex() + symbolDivExpr.getPosition()] = 1;
+      return success();
+    }
     return addLocalVariableSemiAffine(divLhs, rhs, divExpr, lhs, lhs.size());
   }