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

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.

Reviewed By: nicolasvasilache, bondhugula

Differential Revision: https://reviews.llvm.org/D114240

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 783e1be..6faf685 100644 (file)
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -385,8 +385,7 @@ static AffineMap linearizeCollapsedDims(AffineMap sourceMap,
         makeCanonicalStridedLayoutExpr(sizes, dimExprs, context);
     resultExprs.push_back(linearizedExpr);
   }
-  return AffineMap::get(sourceMap.getNumDims(), sourceMap.getNumSymbols(),
-                        resultExprs, context);
+  return AffineMap::inferFromExprList({resultExprs}).front();
 }
 
 // TensorExpandShapeOp is fusable with its consumer (i.e. reshape as a