[MLIR] Avoid creation of buggy affine maps while replacing dimension and symbol

[arnab-oss]

Summary

Initially before appending the newly composed dimension and symbols
to the dimension and symbol list whose size is to be passed in
AffineMap::get(), the call to the AffineMap::get() was made, resulting
in wrong dimCount and symbolCount being passed as argument. We move the
call to the AffineMap::get() after the diimension and symbol list are
updated.

[bondhugula]

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

628

Tweak this comment to be consistent with the change and its rationale.

[nicolasvasilache]

thanks!

[arnab-oss]

Fixed documentation @bondhugula .

[bondhugula]

The relevant commits aren't properly stacked here. D114237 should have been at the bottom but it's instead on top of D114238. Please reorder to fix.

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

static LogicalResult replaceDimOrSym(AffineMap *map,
   assert(composeMap.getNumResults() == 1 && "affine.apply with >1 results");
   AffineExpr composeExpr =
       composeMap.shiftDims(dims.size()).shiftSymbols(syms.size()).getResult(0);
   ValueRange composeDims =
       affineApply.getMapOperands().take_front(composeMap.getNumDims());
   ValueRange composeSyms =
       affineApply.getMapOperands().take_back(composeMap.getNumSymbols());
 
-  // Perform the replacement and append the dims and symbols where relevant.
+  // Append the dims and symbols where relevant and perform the replacement.
   MLIRContext *ctx = map->getContext();
   AffineExpr toReplace = isDimReplacement ? getAffineDimExpr(pos, ctx)
                                           : getAffineSymbolExpr(pos, ctx);
-  *map = map->replace(toReplace, composeExpr, dims.size(), syms.size());
   dims.append(composeDims.begin(), composeDims.end());
   syms.append(composeSyms.begin(), composeSyms.end());
+  *map = map->replace(toReplace, composeExpr, dims.size(), syms.size());
 
   return success();
 }
 
 /// Iterate over `operands` and fold away all those produced by an AffineApplyOp
 /// iteratively. Perform canonicalization of map and operands as well as
 /// AffineMap simplification. `map` and `operands` are mutated in place.
 static void composeAffineMapAndOperands(AffineMap *map,