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

arnab-cerebras · bondhugula · commit e4e4da86aff5 · 2021-11-27T00:32:58.000+05:30
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
@@ -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

Original file line number	Diff line number	Diff line change
`@@ -385,8 +385,7 @@ static AffineMap linearizeCollapsedDims(AffineMap sourceMap,`
`385`	`385`	`makeCanonicalStridedLayoutExpr(sizes, dimExprs, context);`
`386`	`386`	`resultExprs.push_back(linearizedExpr);`
`387`	`387`	`}`
`388`		`- return AffineMap::get(sourceMap.getNumDims(), sourceMap.getNumSymbols(),`
`389`		`- resultExprs, context);`
	`388`	`+ return AffineMap::inferFromExprList({resultExprs}).front();`
`390`	`389`	`}`
`391`	`390`
`392`	`391`	`// TensorExpandShapeOp is fusable with its consumer (i.e. reshape as a`