Cleanup post cl/235283610 - NFC

- remove stale comments + cleanup
- drop MLIRContext * field from expr flattener

PiperOrigin-RevId: 235621178

mlir/include/mlir/IR/AffineExprVisitor.h

@@ -260,15 +260,13 @@ public:
   // will be, and linearize this to std::vector<int64_t> to prevent
   // SmallVector moves on re-allocation.
   std::vector<SmallVector<int64_t, 8>> operandExprStack;
-  // Constraints connecting newly introduced local variables (for mod's and
-  // div's) to existing (dimensional and symbolic) ones. These are always
-  // inequalities.
 
   unsigned numDims;
   unsigned numSymbols;
-  // Number of newly introduced identifiers to flatten mod/floordiv/ceildiv
-  // expressions that could not be simplified.
+
+  // Number of newly introduced identifiers to flatten mod/floordiv/ceildiv's.
   unsigned numLocals;
+
   // AffineExpr's corresponding to the floordiv/ceildiv/mod expressions for
   // which new identifiers were introduced; if the latter do not get canceled
   // out, these expressions can be readily used to reconstruct the AffineExpr
@@ -277,15 +275,19 @@ public:
   // will be simplified to d0 + q, where q = (d0 + d1) ceildiv 2. (d0 + d1)
   // ceildiv 2 would be the local expression stored for q.
   SmallVector<AffineExpr, 4> localExprs;
-  MLIRContext *context;
 
-  SimpleAffineExprFlattener(unsigned numDims, unsigned numSymbols,
-                            MLIRContext *context);
+  SimpleAffineExprFlattener(unsigned numDims, unsigned numSymbols);
 
   virtual ~SimpleAffineExprFlattener() = default;
 
+  // Visitor method overrides.
   void visitMulExpr(AffineBinaryOpExpr expr);
   void visitAddExpr(AffineBinaryOpExpr expr);
+  void visitDimExpr(AffineDimExpr expr);
+  void visitSymbolExpr(AffineSymbolExpr expr);
+  void visitConstantExpr(AffineConstantExpr expr);
+  void visitCeilDivExpr(AffineBinaryOpExpr expr);
+  void visitFloorDivExpr(AffineBinaryOpExpr expr);
 
   //
   // t = expr mod c   <=>  t = expr - c*q and c*q <= expr <= c*q + c - 1
@@ -295,6 +297,16 @@ public:
   // 'expr - c * q' and c * q <= expr <= c * q + c - 1 are added to localVarCst.
   void visitModExpr(AffineBinaryOpExpr expr);
 
+protected:
+  // Add a local identifier (needed to flatten a mod, floordiv, ceildiv expr).
+  // The local identifier added is always a floordiv of a pure add/mul affine
+  // function of other identifiers, coefficients of which are specified in
+  // dividend and with respect to a positive constant divisor. localExpr is the
+  // simplified tree expression (AffineExpr) corresponding to the quantifier.
+  virtual void addLocalFloorDivId(ArrayRef<int64_t> dividend, int64_t divisor,
+                                  AffineExpr localExpr);
+
+private:
   // t = expr floordiv c   <=> t = q, c * q <= expr <= c * q + c - 1
   // A floordiv is thus flattened by introducing a new local variable q, and
   // replacing that expression with 'q' while adding the constraints
@@ -305,21 +317,6 @@ public:
   // t = expr ceildiv c   <=> t =  (expr + c - 1) floordiv c
   void visitDivExpr(AffineBinaryOpExpr expr, bool isCeil);
 
-  void visitDimExpr(AffineDimExpr expr);
-  void visitSymbolExpr(AffineSymbolExpr expr);
-  void visitConstantExpr(AffineConstantExpr expr);
-  void visitCeilDivExpr(AffineBinaryOpExpr expr);
-  void visitFloorDivExpr(AffineBinaryOpExpr expr);
-
-protected:
-  // Add a local identifier (needed to flatten a mod, floordiv, ceildiv expr).
-  // The local identifier added is always a floordiv of a pure add/mul affine
-  // function of other identifiers, coefficients of which are specified in
-  // dividend and with respect to a positive constant divisor. localExpr is the
-  // simplified tree expression (AffineExpr) corresponding to the quantifier.
-  virtual void addLocalFloorDivId(ArrayRef<int64_t> dividend, int64_t divisor,
-                                  AffineExpr localExpr);
-
   int findLocalId(AffineExpr localExpr);
 
   inline unsigned getNumCols() const {

mlir/lib/Analysis/AffineStructures.cpp

@@ -41,7 +41,7 @@ namespace {
 // See comments for SimpleAffineExprFlattener.
 // An AffineExprFlattener extends a SimpleAffineExprFlattener by recording
 // constraint information associated with mod's, floordiv's, and ceildiv's
-// in localVarCst.
+// in FlatAffineConstraints 'localVarCst'.
 struct AffineExprFlattener : public SimpleAffineExprFlattener {
 public:
   // Constraints connecting newly introduced local variables (for mod's and
@@ -50,7 +50,7 @@ public:
   FlatAffineConstraints localVarCst;
 
   AffineExprFlattener(unsigned nDims, unsigned nSymbols, MLIRContext *ctx)
-      : SimpleAffineExprFlattener(nDims, nSymbols, ctx) {
+      : SimpleAffineExprFlattener(nDims, nSymbols) {
     localVarCst.reset(nDims, nSymbols, /*numLocals=*/0);
   }
 
@@ -58,8 +58,9 @@ private:
   // Add a local identifier (needed to flatten a mod, floordiv, ceildiv expr).
   // The local identifier added is always a floordiv of a pure add/mul affine
   // function of other identifiers, coefficients of which are specified in
-  // dividend and with respect to a positive constant divisor. localExpr is the
-  // simplified tree expression (AffineExpr) corresponding to the quantifier.
+  // `dividend' and with respect to the positive constant `divisor'. localExpr
+  // is the simplified tree expression (AffineExpr) corresponding to the
+  // quantifier.
   void addLocalFloorDivId(ArrayRef<int64_t> dividend, int64_t divisor,
                           AffineExpr localExpr) override {
     SimpleAffineExprFlattener::addLocalFloorDivId(dividend, divisor, localExpr);

mlir/lib/IR/AffineExpr.cpp

@@ -336,9 +336,8 @@ AffineExpr mlir::toAffineExpr(ArrayRef<int64_t> eq, unsigned numDims,
 }
 
 SimpleAffineExprFlattener::SimpleAffineExprFlattener(unsigned numDims,
-                                                     unsigned numSymbols,
-                                                     MLIRContext *context)
-    : numDims(numDims), numSymbols(numSymbols), numLocals(0), context(context) {
+                                                     unsigned numSymbols)
+    : numDims(numDims), numSymbols(numSymbols), numLocals(0) {
   operandExprStack.reserve(8);
 }
 
@@ -412,6 +411,7 @@ void SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) {
   int64_t floorDivisor = rhsConst / static_cast<int64_t>(gcd);
 
   // Construct the AffineExpr form of the floordiv to store in localExprs.
+  MLIRContext *context = expr.getContext();
   auto dividendExpr =
       toAffineExpr(floorDividend, numDims, numSymbols, localExprs, context);
   auto divisorExpr = getAffineConstantExpr(floorDivisor, context);
@@ -495,6 +495,7 @@ void SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr,
   // the ceil/floor expr (simplified up until here). Add an existential
   // quantifier to express its result, i.e., expr1 div expr2 is replaced
   // by a new identifier, q.
+  MLIRContext *context = expr.getContext();
   auto a = toAffineExpr(lhs, numDims, numSymbols, localExprs, context);
   auto b = getAffineConstantExpr(divisor, context);
 
@@ -533,7 +534,7 @@ void SimpleAffineExprFlattener::addLocalFloorDivId(ArrayRef<int64_t> dividend,
     subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
   localExprs.push_back(localExpr);
   numLocals++;
-  // dividend and divisor are ignored; an override of this method uses it.
+  // dividend and divisor are not used here; an override of this method uses it.
 }
 
 int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) {
@@ -551,7 +552,7 @@ AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims,
   if (!expr.isPureAffine())
     return expr;
 
-  SimpleAffineExprFlattener flattener(numDims, numSymbols, expr.getContext());
+  SimpleAffineExprFlattener flattener(numDims, numSymbols);
   flattener.walkPostOrder(expr);
   ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
   auto simplifiedExpr = toAffineExpr(flattenedExpr, numDims, numSymbols,
@@ -572,8 +573,7 @@ static bool getFlattenedAffineExprs(
     return true;
   }
 
-  SimpleAffineExprFlattener flattener(numDims, numSymbols,
-                                      exprs[0].getContext());
+  SimpleAffineExprFlattener flattener(numDims, numSymbols);
   // Use the same flattener to simplify each expression successively. This way
   // local identifiers / expressions are shared.
   for (auto expr : exprs) {