Add MLIRContext to concatAffineMaps

Signed-off-by: Ian Wood <[email protected]>

Author: IanWood1

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

@@ -741,7 +741,7 @@ def LinalgStructuredInterface
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         auto maps =  $_op.getIndexingMapsArray();
-        return concatAffineMaps(maps);
+        return concatAffineMaps(maps, $_op.getContext());
       }]
     >,
     InterfaceMethod<

mlir/include/mlir/IR/AffineMap.h

@@ -595,10 +595,7 @@ AffineMap inverseAndBroadcastProjectedPermutation(AffineMap map);
 /// potentially empty maps. Assumes each of the underlying map has 0 symbols.
 /// The resulting map has a number of dims equal to the max of `maps`' dims and
 /// the concatenated results as its results.
-///
-/// This method asserts when `maps` is empty.
-/// TODO: this should return an empty map when `maps` is empty but there is no
-/// way to get the MLIRContext needed to construct it.
+/// Returns an empty map if all input `maps` are empty.
 ///
 /// Example:
 /// When applied to the following list of 3 affine maps,
@@ -616,7 +613,7 @@ AffineMap inverseAndBroadcastProjectedPermutation(AffineMap map);
 /// ```mlir
 ///     (i, j, k) -> (i, k, k, j, i, j)
 /// ```
-AffineMap concatAffineMaps(ArrayRef<AffineMap> maps);
+AffineMap concatAffineMaps(ArrayRef<AffineMap> maps, MLIRContext *context);
 
 /// Returns the map that results from projecting out the dimensions specified in
 /// `projectedDimensions`. The projected dimensions are set to 0.

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

@@ -54,7 +54,8 @@ bool linalg::detail::canOpOperandsBeDroppedImpl(
     // if the op has no loops.
     return linalgOp.getNumLoops() == 0;
   }
-  return inversePermutation(concatAffineMaps(indexingMaps)) != AffineMap();
+  return inversePermutation(concatAffineMaps(
+             indexingMaps, linalgOp.getContext())) != AffineMap();
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -392,7 +392,8 @@ linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
   // 1. Check if any of the iteration dimensions are unit-trip count. They will
   //    end up being unit-trip count if they are used to index into a unit-dim
   //    tensor/memref.
-  AffineMap invertedMap = inversePermutation(concatAffineMaps(indexingMaps));
+  AffineMap invertedMap =
+      inversePermutation(concatAffineMaps(indexingMaps, rewriter.getContext()));
   if (!invertedMap) {
     return rewriter.notifyMatchFailure(genericOp,
                                        "invalid indexing maps for operation");
@@ -486,7 +487,8 @@ linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
   // Abort if the indexing maps of the result operation are not invertible
   // (i.e. not legal) or if no dimension was reduced.
   if (newIndexingMaps == indexingMaps ||
-      !inversePermutation(concatAffineMaps(newIndexingMaps)))
+      !inversePermutation(
+          concatAffineMaps(newIndexingMaps, rewriter.getContext())))
     return failure();
 
   Location loc = genericOp.getLoc();

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

@@ -88,13 +88,18 @@ static bool isOpOperandCanBeDroppedAfterFusedLinalgs(
       indexingMaps.push_back(op.getMatchingIndexingMap(&opOperand));
     }
   }
+  if (indexingMaps.empty()) {
+    // If there are no indexing maps, the operand can only be dropped
+    // if neither ops op have loops.
+    return producer.getNumLoops() == 0 && consumer.getNumLoops() == 0;
+  }
 
   // The concatanation of the remained indexing maps must be invertible, so
   // the bounds of the op can be still computed after dropping the selected
   // operand. inversePermutation returns an empty AffineMap in case the
   // concatanated indexing maps are not invertible.
-  return !indexingMaps.empty() &&
-         inversePermutation(concatAffineMaps(indexingMaps)) != AffineMap();
+  return inversePermutation(concatAffineMaps(
+             indexingMaps, producer.getContext())) != AffineMap();
 }
 
 /// Returns a set of indices of the producer's results which would
@@ -1996,7 +2001,8 @@ class FoldScalarOrSplatConstant : public OpRewritePattern<GenericOp> {
             genericOp.getMatchingIndexingMap(&outputOperand));
 
       // Check if the operation shapes to loops map is computable.
-      if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) {
+      if (!inversePermutation(
+              concatAffineMaps(fusedIndexMaps, rewriter.getContext()))) {
         return rewriter.notifyMatchFailure(
             genericOp, "fused op loop bound computation failed");
       }

mlir/lib/IR/AffineMap.cpp

@@ -833,8 +833,10 @@ AffineMap mlir::inverseAndBroadcastProjectedPermutation(AffineMap map) {
   return AffineMap::get(map.getNumResults(), /*symbolCount=*/0, exprs, context);
 }
 
-AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
-  assert(maps.size());
+AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps,
+                                 MLIRContext *context) {
+  if (maps.empty())
+    return AffineMap::get(context);
   unsigned numResults = 0, numDims = 0, numSymbols = 0;
   for (auto m : maps)
     numResults += m.getNumResults();
@@ -847,8 +849,7 @@ AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
     numSymbols += m.getNumSymbols();
     numDims = std::max(m.getNumDims(), numDims);
   }
-  return AffineMap::get(numDims, numSymbols, results,
-                        maps.front().getContext());
+  return AffineMap::get(numDims, numSymbols, results, context);
 }
 
 /// Common implementation to project out dimensions or symbols from an affine