[MLIR] Prevent creation of buggy affine map after linearizing collapsed dimensions of source map

[arnab-oss]

Summary

Initially we were passing wrong numSymbols argument while calling
AffineMap::get() for creaating affine map with linearized result
expressions. The main problems was the number of symbols of the newly
to be created map may be different from that of the source map, as
new symbolic identifiers may be introduced while creating strided layout
linearized expressions.

[bondhugula]

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

390–410

You can completely get rid of hasEncounteredSymbol if you just actually compute numSymbols. You can drop maxSymbolPosition as well.

numSymbols = std::max(numSymbols, symbolExpr.getPosition() + 1);

A single variable will do and it's more natural and compact.

[bondhugula]

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

390–410

Just initialize numSymbols = 0 on the top and it'll take care of everything. No additional conditional needed as well.

411

numSymbol -> numSymbols

[arnab-oss]

Addressed comments by @bondhugula.

[bondhugula]

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

389

Nit: sourceMap in backticks.

407

I assume the existing test cases already exercise this. Update a test case to test this?

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

390

This revision is also duplicating existing functionality, and implementing it with recursion, just like in D114238.

Either: return AffineMap::inferFromtExprList(...);

Absolute worst case, you can expose:

template <typename AffineExprContainer>
static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList,
                               int64_t &maxDim, int64_t &maxSym) {

as a util and reuse it rather than reimplement a 5 line iterative algorithm with a longer recursive one.

[arnab-oss]

Addressed comments by @nicolasvasilache and @bondhugula.

[arnab-oss]

ping @nicolasvasilache, I have addressed your comments. You can take a look.

[nicolasvasilache]

thanks for the fix!

[bondhugula]

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

388–412

Do you need the extra parentheses?

[arnab-oss]

Removed extra parenthesis as suggested by @bondhugula.

[bondhugula]

This revision isn't properly based -- please rebase on the current master tip and update.

[arnab-oss]

Rebased on latest master.

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

for (auto en : llvm::zip(sourceShape.slice(indices[0], indices.size()),
         continue;
       sizes.push_back(std::get<0>(en));
       dimExprs.push_back(std::get<1>(en));
     }
     AffineExpr linearizedExpr =
         makeCanonicalStridedLayoutExpr(sizes, dimExprs, context);
     resultExprs.push_back(linearizedExpr);
   }
-  return AffineMap::get(sourceMap.getNumDims(), sourceMap.getNumSymbols(),
-                        resultExprs, context);
+  // Number of symbols may change as compared to what is present in the
+  // sourceMap as stride length may be constant or symbolic.
+  unsigned maxSymbolPosition = 0;
+  bool hasEncounteredSymbol = false;
+  std::function<void(AffineExpr)> updateMaxSymbolPosition =
+      [&](AffineExpr expr) {
+        auto binOpExpr = expr.dyn_cast<AffineBinaryOpExpr>();
+        if (binOpExpr) {
+          updateMaxSymbolPosition(binOpExpr.getLHS());
+          updateMaxSymbolPosition(binOpExpr.getRHS());
+          return;
+        }
+        if (auto symbolExpr = expr.dyn_cast<AffineSymbolExpr>()) {
+          maxSymbolPosition =
+              std::max(maxSymbolPosition, symbolExpr.getPosition());
+          hasEncounteredSymbol = true;
+        }
+      };
+
+  for (AffineExpr expr : resultExprs)
+    updateMaxSymbolPosition(expr);
+
+  unsigned numSymbol = hasEncounteredSymbol ? 1 + maxSymbolPosition : 0;
+  return AffineMap::get(sourceMap.getNumDims(), numSymbol, resultExprs,
+                        context);
 }
 
 // TensorExpandShapeOp is fusable with its consumer (i.e. reshape as a
 // producer). Fusing when operand has higher rank will require use of mods and
 // divs in the indexing maps of the fused op which would make it non-invertible.
 static bool isTensorReshapeOpFoldableByLinearization(
     TensorExpandShapeOp expandOp, AffineMap useIndexMap, bool asProducer) {
   if (!asProducer)