[mlir][Affine] Fix bug in Affine canonicalization

[philass]

Summary

Both lhs, and rhs need to be updated after binExpr has been updated.

The call to getAffineBinaryOpExpr, can call a function like simplifyFloorDiv. This function can change not only the binaryOp, but also the lhs and rhs of the AffineBinaryOp being constructed.

[ftynse]

Please add a test

[bondhugula]

I actually have this fix with a test case. I can send a patch.

[bondhugula]

This is addressed by https://reviews.llvm.org/D138584.

[philass]

Thanks bondhugula

I'll close this.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

static void simplifyExprAndOperands(AffineExpr &expr,
   expr = getAffineBinaryOpExpr(binExpr.getKind(), lhs, rhs);
 
   binExpr = expr.dyn_cast<AffineBinaryOpExpr>();
   if (!binExpr || (binExpr.getKind() != AffineExprKind::FloorDiv &&
                    binExpr.getKind() != AffineExprKind::Mod)) {
     return;
   }
 
+  lhs = binExpr.getLHS();
+  rhs = binExpr.getRHS();
+
   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
   if (!rhsConst)
     return;
 
   int64_t rhsConstVal = rhsConst.getValue();
   AffineExpr quotientTimesDiv, rem;
   int64_t divisor;