[mlir][IR] Add optional offset `offset` parameter to shiftDims/shiftSymbols

[springerm]

Summary

Dimensions/symbols are shifted, starting from the specified offset. This allows for inserting new dims/syms into an AffineMap at an arbitrary position (not just at the beginning).

Depends On D108185

[ftynse]

It would be nice to fold this into the commit that uses the new functionality so that we don't go with untested code for a long time.

[springerm]

rebase

[springerm]

rebase

[bondhugula]

mlir/include/mlir/IR/AffineMap.h

210–222

This doc comment is now a bit convoluted. Can you please rephrase it simply:

/// Replace dims[offset ... numDims) by dims[offset + shift ... ...

mlir/lib/IR/AffineExpr.cpp

116–121

Likewise.

[springerm]

address comments

mlir/include/mlir/IR/AffineMap.h

210–222

How about this?

/// Replace dims[offset ... numDims)
/// by dims[offset + shift ... shift + numDims).

[nicolasvasilache]

Thanks for extracting this to make it super easy to review.
Please fold into the user when landing as per @ftynse 's comment.

mlir/include/mlir/IR/AffineExpr.h

154

nit: can we align dims[ across lines?

[springerm]

rebase

[springerm]

Merged into and committed with D107222

mlir/include/mlir/IR/AffineExpr.h

public:
   /// modified expression tree.
   AffineExpr replace(AffineExpr expr, AffineExpr replacement) const;
 
   /// Sparse replace method. If `*this` appears in `map` replaces it by
   /// `map[*this]` and return the modified expression tree. Otherwise traverse
   /// `*this` and apply replace with `map` on its subexpressions.
   AffineExpr replace(const DenseMap<AffineExpr, AffineExpr> &map) const;
 
-  /// Replace dims[0 .. numDims - 1] by dims[shift .. shift + numDims - 1].
-  AffineExpr shiftDims(unsigned numDims, unsigned shift) const;
+  /// Replace dims[0 .. numDims - 1] by
+  ///         dims[0 .. offset - 1] +
+  ///         dims[offset + shift .. shift + numDims - 1].
+  AffineExpr shiftDims(unsigned numDims, unsigned shift,
+                       unsigned offset = 0) const;
 
   /// Replace symbols[0 .. numSymbols - 1] by
-  ///         symbols[shift .. shift + numSymbols - 1].
-  AffineExpr shiftSymbols(unsigned numSymbols, unsigned shift) const;
+  ///         symbols[0 .. offset - 1] +
+  ///         symbols[offset + shift .. shift + numSymbols - 1].
+  AffineExpr shiftSymbols(unsigned numSymbols, unsigned shift,
+                          unsigned offset = 0) const;
 
   AffineExpr operator+(int64_t v) const;
   AffineExpr operator+(AffineExpr other) const;
   AffineExpr operator-() const;
   AffineExpr operator-(int64_t v) const;
   AffineExpr operator-(AffineExpr other) const;
   AffineExpr operator*(int64_t v) const;
   AffineExpr operator*(AffineExpr other) const;

mlir/include/mlir/IR/AffineMap.h

public:
   AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map) const;
 
   /// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the
   /// results and return a new AffineMap with the new results and with the
   /// specified number of dims and symbols.
   AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map,
                     unsigned numResultDims, unsigned numResultSyms) const;
 
-  /// Replace dims[0 .. numDims - 1] by dims[shift .. shift + numDims - 1].
-  AffineMap shiftDims(unsigned shift) const {
-    return AffineMap::get(
-        getNumDims() + shift, getNumSymbols(),
+  /// Replace dims[0 .. numDims - 1] by
+  ///         dims[0 .. offset - 1] +
+  ///         dims[offset + shift .. shift + numDims - 1].
+  AffineMap shiftDims(unsigned shift, unsigned offset = 0) const {
+    assert(offset <= getNumDims());
+    return AffineMap::get(getNumDims() + shift, getNumSymbols(),
                           llvm::to_vector<4>(llvm::map_range(
                               getResults(),
-            [&](AffineExpr e) { return e.shiftDims(getNumDims(), shift); })),
+                              [&](AffineExpr e) {
+                                return e.shiftDims(getNumDims(), shift, offset);
+                              })),
                           getContext());
   }
 
   /// Replace symbols[0 .. numSymbols - 1] by
-  ///         symbols[shift .. shift + numSymbols - 1].
-  AffineMap shiftSymbols(unsigned shift) const {
+  ///         symbols[0 .. offset - 1] +
+  ///         symbols[offset + shift .. shift + numSymbols - 1].
+  AffineMap shiftSymbols(unsigned shift, unsigned offset = 0) const {
     return AffineMap::get(getNumDims(), getNumSymbols() + shift,
                           llvm::to_vector<4>(llvm::map_range(
                               getResults(),
                               [&](AffineExpr e) {
-                                return e.shiftSymbols(getNumSymbols(), shift);
+                                return e.shiftSymbols(getNumSymbols(), shift,
+                                                      offset);
                               })),
                           getContext());
   }
 
   /// Folds the results of the application of an affine map on the provided
   /// operands to a constant if possible.
   LogicalResult constantFold(ArrayRef<Attribute> operandConstants,
                              SmallVectorImpl<Attribute> &results) const;

mlir/lib/IR/AffineExpr.cpp

AffineExpr AffineExpr::replaceDims(ArrayRef<AffineExpr> dimReplacements) const {
   return replaceDimsAndSymbols(dimReplacements, {});
 }
 
 AffineExpr
 AffineExpr::replaceSymbols(ArrayRef<AffineExpr> symReplacements) const {
   return replaceDimsAndSymbols({}, symReplacements);
 }
 
-/// Replace symbols[0 .. numDims - 1] by symbols[shift .. shift + numDims - 1].
-AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift) const {
+/// Replace dims[0 .. numDims - 1] by
+///         dims[0 .. offset - 1] + dims[offset + shift .. shift + numDims - 1].
+AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift,
+                                 unsigned offset) const {
   SmallVector<AffineExpr, 4> dims;
-  for (unsigned idx = 0; idx < numDims; ++idx)
+  for (unsigned idx = 0; idx < offset; ++idx)
+    dims.push_back(getAffineDimExpr(idx, getContext()));
+  for (unsigned idx = offset; idx < numDims; ++idx)
     dims.push_back(getAffineDimExpr(idx + shift, getContext()));
   return replaceDimsAndSymbols(dims, {});
 }
 
 /// Replace symbols[0 .. numSymbols - 1] by
-/// symbols[shift .. shift + numSymbols - 1].
-AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift) const {
+///         symbols[0 .. offset - 1] +
+///         symbols[offset + shift .. shift + numSymbols - 1].
+AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift,
+                                    unsigned offset) const {
   SmallVector<AffineExpr, 4> symbols;
-  for (unsigned idx = 0; idx < numSymbols; ++idx)
+  for (unsigned idx = 0; idx < offset; ++idx)
+    symbols.push_back(getAffineSymbolExpr(idx, getContext()));
+  for (unsigned idx = offset; idx < numSymbols; ++idx)
     symbols.push_back(getAffineSymbolExpr(idx + shift, getContext()));
   return replaceDimsAndSymbols({}, symbols);
 }
 
 /// Sparse replace method. Return the modified expression tree.
 AffineExpr
 AffineExpr::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {
   auto it = map.find(*this);