[mlir][Affine] Add `MLIRContext` argument to `concatAffineMaps`.

[mravishankar]

Summary

Avoid segfault in concatAffineMaps, when the first map of the list
of maps to be concatenated is querried for the context. If no maps are
passed then this leads to a segfault. Pass the context explicitly and
return an empty map instead. This has a cascading effect where
context needs to be passed into many methods.

Depends D138096

[nicolasvasilache]

Generally looks good, thanks Mahesh!
I flagged one place where it seems things could be simplified a bit

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

37

Could we use the linalgOp to derive the context here and avoid passing it through this path and callers?

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

 class IteratorTypeAttr;
 class LinalgOp;
 
 namespace detail {
 /// Implementation of the method that that check if given operands
 /// can be dropped, i.e. the remaining operands can compute the loop
 /// bounds of the op.
 bool canOpOperandsBeDroppedImpl(linalg::LinalgOp linalgOp,
-                                ArrayRef<OpOperand *> droppedOperands);
+                                ArrayRef<OpOperand *> droppedOperands,
+                                MLIRContext *context);
 } // namespace detail
 
 /// Checks whether `linalgOp` conforms to ContractionOpInterface.
 // TODO: embed within `isa<ContractionOpInterface>` if possible / natural.
 bool isaContractionOpInterface(LinalgOp linalgOp);
 
 namespace detail {
 

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

InterfaceMethod<
         compute subshapes even in cases where the concatenation of indexing maps
         (i.e. the data traversal order) is not a simple permutation of the loop
         traversal order. It is then possible to define ops with skewed data
         traversal order for which we can still easily compute hyperrectangular
         loop bounds and subviews.
       }],
       /*retTy=*/"AffineMap",
       /*methodName=*/"getLoopsToShapesMap",
-      /*args=*/(ins),
+      /*args=*/(ins "MLIRContext *":$context),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         auto maps =  $_op.getIndexingMapsArray();
-        return concatAffineMaps(maps);
+        return concatAffineMaps(maps, context);
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Hook to provide a custom AffineMap used to construct the
         hyperrectangular loop iteration space given all the operand subshapes.
         This is used to answer the question:
         "Given a list of operand ranges, what is the subportion of the iteration
         space involved in the computation".
         This is the inverse problem of `getLoopsToShapesMap`.
         Return the empty AffineMap when such an AffineMap cannot be constructed.
         The default behavior is based on a very simple inference procedure that
         only works with permutation affine maps.
         A more advanced Tensor-Comprehension like inference is possible but has
         proven to be ambiguous in unfavorable case.
         A safer and more robust alternative is to allow each op to define
         its own AffineMap.
       }],
       /*retTy=*/"AffineMap",
       /*methodName=*/"getShapesToLoopsMap",
-      /*args=*/(ins),
+      /*args=*/(ins "MLIRContext *":$context),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return inversePermutation(getLoopsToShapesMap());
+        return inversePermutation(getLoopsToShapesMap(context));
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Checks if the given operands can be dropped, and the remaining
         operands can still compute the bounds of the op.
       }],
       /*retTy=*/"bool",
       /*methodName=*/"canOpOperandsBeDropped",
-      /*args=*/(ins "ArrayRef<OpOperand *>":$droppedOperands),
+      /*args=*/(ins "MLIRContext *":$context,
+          "ArrayRef<OpOperand *>":$droppedOperands),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return detail::canOpOperandsBeDroppedImpl($_op, droppedOperands);
+        return detail::canOpOperandsBeDroppedImpl($_op, droppedOperands, context);
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Like `getShape`, but only returns statically-known information, without
         generating any new IR. For each shape dimension, returns >=0 if that
         dimension is statically known, or ShapeType::kDynamicSize otherwise.
       }],
InterfaceMethod<
       /*desc=*/[{
         Returns the statically-known loop ranges. Composes
         `getShapesToLoopsMap()` with the result of `getStaticShape`.
         Returns ShapeType::kDynamicSize for non-statically-known loop ranges.
         This is expected to be called by a valid Linalg op
       }],
       /*retTy=*/"SmallVector<int64_t, 4>",
       /*methodName=*/"getStaticLoopRanges",
-      /*args=*/(ins),
+      /*args=*/(ins "MLIRContext *":$context),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         SmallVector<int64_t> viewSizes = getStaticShape();
-        AffineMap invertedMap = getShapesToLoopsMap();
+        AffineMap invertedMap = getShapesToLoopsMap(context);
         assert(invertedMap && "expected a valid Linalg op to call the method");
         return invertedMap.compose(viewSizes);
       }]
     >,
     //===------------------------------------------------------------------===//
     // Other static interface methods.
     //===------------------------------------------------------------------===//
     StaticInterfaceMethod<
let extraClassDeclaration = [{
     /// Create the loop ranges to materialize the computation over the current
     /// operands. This is done by applying `getShapesToLoopsMap` to
     /// `createFlatListOfOperandDims`.
     SmallVector<Range, 4> createLoopRanges(OpBuilder &b, Location loc);
 
     /// Compute the static loop sizes necessary to vectorize the computation.
     /// This is done by applying `getShapesToLoopsMap` to
     /// `createFlatListOfOperandStaticDims`.
-    SmallVector<int64_t, 4> computeStaticLoopSizes();
+    SmallVector<int64_t, 4> computeStaticLoopSizes(MLIRContext *context);
 
     /// Returns the value that expresses the shape of the output in terms of
     /// shape of the input operands where possible
     LogicalResult reifyResultShapes(OpBuilder &b,
         ReifiedRankedShapedTypeDims &reifiedReturnShapes);
 
     //========================================================================//
     // Forwarding functions to access interface methods from the

mlir/include/mlir/IR/AffineMap.h

 ///    }
 /// ```
 ///
 /// Returns the 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.
 ///
 /// Example:
 /// 1) map                  : affine_map<(d0, d1, d2) -> (d0, d1)>
 ///    projected_dimensions : {2}
 ///    result               : affine_map<(d0, d1) -> (d0, d1)>

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

 
 /// Include the definitions of the copy operation interface.
 #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // Interface utility functions
 //===----------------------------------------------------------------------===//
 bool linalg::detail::canOpOperandsBeDroppedImpl(
-    linalg::LinalgOp linalgOp, ArrayRef<OpOperand *> droppedOperands) {
+    linalg::LinalgOp linalgOp, ArrayRef<OpOperand *> droppedOperands,
+    MLIRContext *context) {
   SmallVector<AffineMap> indexingMaps;
   for (auto &opOperand : linalgOp->getOpOperands()) {
     if (llvm::is_contained(droppedOperands, &opOperand))
       continue;
     indexingMaps.push_back(linalgOp.getMatchingIndexingMap(&opOperand));
   }
   if (indexingMaps.empty()) {
     // If there are no indexing maps, the operand can only be dropped
     // if the op has no loops.
     return linalgOp.getNumLoops() == 0;
   }
-  return inversePermutation(concatAffineMaps(indexingMaps)) != AffineMap();
+  return inversePermutation(concatAffineMaps(indexingMaps, context)) !=
+         AffineMap();
 }
 
 //===----------------------------------------------------------------------===//
 // ContractionOpInterface implementation
 //===----------------------------------------------------------------------===//
 
 /// Return true if the use-def chain from `v` to `from` consists of 0 or more
 /// unary single-operand operations.
SmallVector<int64_t, 4> LinalgOp::createFlatListOfOperandStaticDims() {
   SmallVector<int64_t, 4> res;
   assert(!hasDynamicShape() && "expected operands to have static shapes");
   for (OpOperand &opOperand : getOperation()->getOpOperands())
     llvm::append_range(res, getShape(&opOperand));
   return res;
 }
 
 SmallVector<Range, 4> LinalgOp::createLoopRanges(OpBuilder &b, Location loc) {
-  AffineMap map = getLoopsToShapesMap();
+  AffineMap map = getLoopsToShapesMap(b.getContext());
   unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
   auto viewSizes = createFlatListOfOperandDims(b, loc);
   SmallVector<Range, 4> res(numDims);
   for (unsigned idx = 0; idx < numRes; ++idx) {
     auto result = map.getResult(idx);
     if (auto d = result.dyn_cast<AffineDimExpr>()) {
       if (res[d.getPosition()].offset)
         continue;
       res[d.getPosition()] =
           Range{b.getIndexAttr(0), viewSizes[idx], b.getIndexAttr(1)};
     }
   }
   return res;
 }
 
-SmallVector<int64_t, 4> LinalgOp::computeStaticLoopSizes() {
-  AffineMap map = getLoopsToShapesMap();
+SmallVector<int64_t, 4> LinalgOp::computeStaticLoopSizes(MLIRContext *context) {
+  AffineMap map = getLoopsToShapesMap(context);
   unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
   SmallVector<int64_t, 4> allShapeSizes = createFlatListOfOperandStaticDims();
   SmallVector<int64_t, 4> res(numDims, 0);
   for (unsigned idx = 0; idx < numRes; ++idx) {
     auto result = map.getResult(idx);
     if (auto d = result.dyn_cast<AffineDimExpr>())
       res[d.getPosition()] = allShapeSizes[idx];
   }
LinalgOp::reifyResultShapes(OpBuilder &b,
   // Consider the following expression O(i+j, j) += A(i,k) * B(k, j)
   // We want to express the shape of dim 0 of O in terms of shape of the inputs.
   // This is achieved as follows.
   //   loopsToShapesMap = (d0, d1, d2) -> (d0, d2, d2, d1, d0 + d1, d1)
   //   subMapOfResultShapes = (d0, d1, d2) -> (d0 + d1, d1)
   //   shapesToLoopsMap = (d0, d2, d2, d3, d4, d5) -> (d0, d3, d2)
   //   resultShapesFromInputShapes = subMapOfResultDim.compose(shapesToLoopMap)
   //     = (d0, d1, d2, d3, d4, d5) -> (d0 + d1, d1)
-  AffineMap loopsToShapesMap = getLoopsToShapesMap();
+  AffineMap loopsToShapesMap = getLoopsToShapesMap(b.getContext());
 
   // Find the position in the above map that represents the shape of the
   // result:dim being inferred.
   auto resultShapesSubMapPos = getResultsPositionInLoopsToShapeMap(*this);
 
   /// From loopsToShapesMap extract the submap that represents the shape of the
   /// (resultIdx, dim) needed.
   AffineMap loopToResultsShapeMap = loopsToShapesMap.getSliceMap(
       resultShapesSubMapPos.first,
       resultShapesSubMapPos.second - resultShapesSubMapPos.first);
   AffineMap resultShapesFromInputShapesMap =
-      loopToResultsShapeMap.compose(getShapesToLoopsMap());
+      loopToResultsShapeMap.compose(getShapesToLoopsMap(b.getContext()));
 
   // Check that the result dim map does not contain the positions corresponding
   // to the outputs.
   llvm::SmallBitVector outputDims(resultShapesFromInputShapesMap.getNumDims());
   outputDims.set(resultShapesSubMapPos.first, resultShapesSubMapPos.second);
   HasAffineDimExprVisitor checkDimExpr(std::move(outputDims));
   Location loc = getOperation()->getLoc();
   IRRewriter rewriter(b);
if (indexingMap.getNumResults() != rank)
              << rank << ") to match the result rank of indexing_map #"
              << opOperand.getOperandNumber() << " ("
              << indexingMap.getNumResults() << ")";
   }
 
   SmallVector<unsigned> redDims;
   linalgOp.getReductionDims(redDims);
 
-  if (!linalgOp.getShapesToLoopsMap())
+  if (!linalgOp.getShapesToLoopsMap(op->getContext()))
     return op->emitOpError("expected the shape-to-loops map to be non-null");
 
   // Check if given shapes match to inferred shapes.
-  SmallVector<int64_t, 4> endLoopRangeValues = linalgOp.getStaticLoopRanges();
+  SmallVector<int64_t, 4> endLoopRangeValues =
+      linalgOp.getStaticLoopRanges(op->getContext());
   SmallVector<int64_t, 4> startLoopRangeValues(endLoopRangeValues.size(), 0);
 
   // Verify only static cases since we can't get exact dimension sizes and loop
   // ranges for dynamic cases in this stage.
   if (llvm::none_of(endLoopRangeValues, ShapedType::isDynamic)) {
     for (int64_t &range : endLoopRangeValues)
       range -= 1;
     for (OpOperand &opOperand : linalgOp->getOpOperands()) {

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

transform::ScalarizeOp::applyToOne(linalg::LinalgOp target,
                                    SmallVectorImpl<Operation *> &results,
                                    transform::TransformState &state) {
   scf::SCFTilingOptions tilingOptions;
   tilingOptions.setTileSizeComputationFunction([&](OpBuilder &b, Operation *) {
     SmallVector<Value, 4> tileSizes;
     Location loc = target.getLoc();
     SmallVector<OpFoldResult> allShapeSizes =
         target.createFlatListOfOperandDims(b, loc);
-    AffineMap map = target.getShapesToLoopsMap();
+    AffineMap map = target.getShapesToLoopsMap(getContext());
     if (!map)
       return tileSizes;
     IRRewriter rewriter(b);
     SmallVector<OpFoldResult> shapeSizes =
         makeComposedFoldedMultiResultAffineApply(rewriter, loc, map,
                                                  allShapeSizes);
     // If the shape size is dynamic, tile by 1.
     // Otherwise, do not tile (i.e. tile size 0).

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

LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
     if (!indexingMap.isProjectedPermutation()) {
       return rewriter.notifyMatchFailure(
           sliceOp, "expected a projected permutation for output");
     }
 
     auto linalgLoc = linalgOp.getLoc();
     SmallVector<OpFoldResult> allShapeSizes =
         linalgOp.createFlatListOfOperandDims(rewriter, linalgLoc);
-    AffineMap shapeSizesToLoopsMap = linalgOp.getShapesToLoopsMap();
+    AffineMap shapeSizesToLoopsMap =
+        linalgOp.getShapesToLoopsMap(rewriter.getContext());
     if (!shapeSizesToLoopsMap) {
       return rewriter.notifyMatchFailure(
           linalgOp, "failed to get loops map from shape sizes");
     }
     SmallVector<OpFoldResult> sizeBounds =
         makeComposedFoldedMultiResultAffineApply(
             rewriter, linalgLoc, shapeSizesToLoopsMap, allShapeSizes);
 

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

LogicalResult matchAndRewrite(GenericOp genericOp,
     // Identified this as a potential candidate for folding. Now check the
     // policy to see whether we are allowed to proceed.
     for (OpOperand *operand : genericOp.getDpsInputOperands()) {
       if (!controlFn(operand))
         return failure();
     }
 
     auto linalgOp = cast<LinalgOp>(genericOp.getOperation());
-    SmallVector<int64_t, 4> loopBounds = linalgOp.computeStaticLoopSizes();
+    SmallVector<int64_t, 4> loopBounds =
+        linalgOp.computeStaticLoopSizes(rewriter.getContext());
     int64_t numElements = outputType.getNumElements();
 
     // Use APInt/APFloat instead of Attribute here for constructing the output.
     // This helps to avoid blowing up compiler memory usage: Attributes would
     // unify the following cases but they have lifetime as the MLIRContext.
     SmallVector<APInt> intOutputValues;
     SmallVector<APFloat> fpOutputValues;
     if (elementType.template isa<FloatType>())

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

 /// Helper method to compute the range of a generic op.
 static SmallVector<OpFoldResult> getGenericOpLoopRange(OpBuilder &b,
                                                        GenericOp op) {
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(op);
   Location loc = op.getLoc();
   auto allShapesSizes =
       cast<LinalgOp>(op.getOperation()).createFlatListOfOperandDims(b, loc);
-  AffineMap map = op.getShapesToLoopsMap();
+  AffineMap map = op.getShapesToLoopsMap(b.getContext());
   IRRewriter rewriter(b);
   return makeComposedFoldedMultiResultAffineApply(rewriter, loc, map,
                                                   allShapesSizes);
 }
 
 /// Helper method to permute the list of `values` based on the `map`.
 SmallVector<OpFoldResult> permuteValues(ArrayRef<OpFoldResult> values,
                                         AffineMap map) {

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

for (AffineMap operandMap : indexingMaps) {
     newIndexingMaps.push_back(simplifyAffineMap(
         operandMap.replaceDimsAndSymbols(dimReplacements, symReplacements,
                                          numIterationDims - unitDims.size(),
                                          numSymbols)));
   }
 
   // Check that the new index maps are invertible. If not, something went
   // wrong, so abort.
-  if (!inversePermutation(concatAffineMaps(newIndexingMaps)))
+  if (!inversePermutation(concatAffineMaps(newIndexingMaps, context)))
     return nullptr;
   return ArrayAttr::get(context,
                         llvm::to_vector<4>(llvm::map_range(
                             newIndexingMaps, [](AffineMap map) -> Attribute {
                               return AffineMapAttr::get(map);
                             })));
 }
 
LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     SmallVector<AffineMap, 4> indexingMaps = genericOp.getIndexingMapsArray();
     if (indexingMaps.empty())
       return failure();
 
     // 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 failure();
     SmallVector<int64_t> dims = genericOp.getStaticShape();
 
     DenseSet<unsigned> unitDims;
     SmallVector<unsigned, 4> unitDimsReductionLoops;
     ArrayAttr iteratorTypes = genericOp.getIteratorTypes();
     for (const auto &expr : enumerate(invertedMap.getResults())) {
for (OpOperand &opOperand : genericOp->getOpOperands()) {
                                  getAffineDimExpr(dim, context), context));
             }));
         reassociationMaps.push_back(ArrayAttr::get(context, maps));
       }
     }
 
     // If the indexing maps of the result operation are not invertible (i.e. not
     // legal), abort.
-    if (!doCanonicalization ||
-        !inversePermutation(concatAffineMaps(newIndexingMaps)))
+    if (!doCanonicalization || !inversePermutation(concatAffineMaps(
+                                   newIndexingMaps, rewriter.getContext())))
       return failure();
 
     // If any operand type change, insert a reshape to convert from the original
     // type to the new type.
     // TODO: get rid of flattenedIdx which assumes operand order and contiguity.
     unsigned flattenedIdx = 0;
     auto insertReshapes = [&](ValueRange values) {
       SmallVector<Value, 4> res;

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

mlir::linalg::fuseElementwiseOps(RewriterBase &rewriter,
   // TODO: allow fusing the producer of an output operand.
   assert(consumer.isDpsInput(fusedOperand) &&
          "expected producer of input operand");
   /// Find the results of the producer that have uses outside of the consumer.
   llvm::SmallDenseSet<int> preservedProducerResults;
   for (auto producerResult : llvm::enumerate(producer->getResults())) {
     auto outputOperand = producer.getDpsInitOperand(producerResult.index());
     if (producer.payloadUsesValueFromOperand(outputOperand) ||
-        !producer.canOpOperandsBeDropped(outputOperand) ||
+        !producer.canOpOperandsBeDropped(rewriter.getContext(),
+                                         outputOperand) ||
         llvm::any_of(producerResult.value().getUsers(), [&](Operation *user) {
           return user != consumer.getOperation();
         })) {
       preservedProducerResults.insert(producerResult.index());
     }
   }
 
   // Compute the fused operands list and indexing maps.
mlir::linalg::fuseElementwiseOps(RewriterBase &rewriter,
 
   // Generate the fused op.
   auto fusedOp = rewriter.create<GenericOp>(
       consumer.getLoc(), fusedResultTypes, fusedInputOperands,
       fusedOutputOperands, rewriter.getAffineMapArrayAttr(fusedIndexMaps),
       consumer.getIteratorTypes(),
       /*doc=*/nullptr,
       /*library_call=*/nullptr);
-  if (!fusedOp.getShapesToLoopsMap()) {
+  if (!fusedOp.getShapesToLoopsMap(rewriter.getContext())) {
     // Fused op has invalid indexing maps. Typically this means something is off
     // in the input, but going ahead here would result in verification errors.
     // So cleanup and abort.
     rewriter.eraseOp(fusedOp);
     return rewriter.notifyMatchFailure(
         fusedOp, "fused op failed loop bound computation check");
   }
 
LogicalResult ExpansionInfo::compute(LinalgOp linalgOp,
                                      ArrayRef<AffineMap> reassociationMaps,
                                      ArrayRef<int64_t> expandedShape,
                                      ArrayRef<int64_t> collapsedShape,
                                      PatternRewriter &rewriter) {
   if (reassociationMaps.empty())
     return failure();
   AffineMap fusedIndexMap = linalgOp.getMatchingIndexingMap(fusableOpOperand);
 
-  SmallVector<int64_t, 4> originalLoopRange = linalgOp.getStaticLoopRanges();
+  SmallVector<int64_t, 4> originalLoopRange =
+      linalgOp.getStaticLoopRanges(rewriter.getContext());
   originalLoopExtent.assign(originalLoopRange.begin(), originalLoopRange.end());
 
   reassociation.clear();
   expandedShapeMap.clear();
   // Compute the number of dimension in the expanded op that correspond to each
   // dimension of the original op.
   SmallVector<unsigned> numExpandedDims(fusedIndexMap.getNumDims(), 1);
   expandedShapeMap.resize(fusedIndexMap.getNumDims());
for (OpOperand *opOperand : genericOp.getDpsInputOperands()) {
         fusedOperands.push_back(inputValue);
         fusedLocs.push_back(inputValue.getLoc());
       }
       for (OpOperand *outputOperand : genericOp.getDpsInitOperands())
         fusedIndexMaps.push_back(
             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");
       }
 
       // Create a constant scalar value from the splat constant.
       Value scalarConstant = rewriter.create<arith::ConstantOp>(
           def->getLoc(), constantAttr, constantAttr.getType());
 

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

LogicalResult matchAndRewrite(GenericOp genericOp,
     SmallVector<OpOperand *> droppedOpOperands;
 
     // Information needed to build the new op.
     SmallVector<Value> newInputOperands, newOutputOperands;
     SmallVector<AffineMap> newIndexingMaps;
 
     // Gather information about duplicate input operands.
     llvm::SmallDenseMap<unsigned, unsigned> origInsToNewInsPos =
-        deduplicateInputOperands(genericOp, droppedOpOperands, newInputOperands,
+        deduplicateInputOperands(rewriter.getContext(), genericOp,
+                                 droppedOpOperands, newInputOperands,
                                  newIndexingMaps);
 
     // Gather information about the dropped outputs.
     llvm::SmallDenseMap<unsigned, unsigned> origOutsToNewOutsPos =
-        deduplicateOutputOperands(genericOp, droppedOpOperands,
-                                  newOutputOperands, newIndexingMaps);
+        deduplicateOutputOperands(rewriter.getContext(), genericOp,
+                                  droppedOpOperands, newOutputOperands,
+                                  newIndexingMaps);
 
     // Check if there is any change to operands.
     if (newInputOperands.size() + newOutputOperands.size() ==
         genericOp->getNumOperands())
       return failure();
 
     // Create the new op with the body being empty.
     Location loc = genericOp.getLoc();
struct DeduplicateAndRemoveDeadOperandsAndResults
   }
 
 private:
   // Deduplicate input operands, and return the
   // - Mapping from operand position in the original op, to operand position in
   // the canonicalized op.
   // - The preserved input operands list (by reference).
   llvm::SmallDenseMap<unsigned, unsigned>
-  deduplicateInputOperands(GenericOp genericOp,
+  deduplicateInputOperands(MLIRContext *context, GenericOp genericOp,
                            SmallVector<OpOperand *> &droppedOpOperands,
                            SmallVector<Value> &newInputOperands,
                            SmallVector<AffineMap> &newIndexingMaps) const {
     llvm::SmallDenseMap<unsigned, unsigned> origToNewPos;
     llvm::SmallDenseMap<std::pair<Value, AffineMap>, unsigned> dedupedInputs;
     for (const auto &en : llvm::enumerate(genericOp.getDpsInputOperands())) {
       OpOperand *inputOpOperand = en.value();
       // Check if operand is dead and if dropping the indexing map makes the
       // loops to shape computation invalid.
       if (!genericOp.payloadUsesValueFromOperand(inputOpOperand)) {
         // Add the current operands to the list of potentially droppable
         // operands. If it cannot be dropped, this needs to be popped back.
         droppedOpOperands.push_back(inputOpOperand);
-        if (genericOp.canOpOperandsBeDropped(droppedOpOperands))
+        if (genericOp.canOpOperandsBeDropped(context, droppedOpOperands))
           continue;
         droppedOpOperands.pop_back();
       }
 
       // Check if this operand is a duplicate.
       AffineMap indexingMap = genericOp.getMatchingIndexingMap(inputOpOperand);
       auto it = dedupedInputs.find(
           std::make_pair(inputOpOperand->get(), indexingMap));
deduplicateInputOperands(MLIRContext *context, GenericOp genericOp,
     return origToNewPos;
   }
 
   // Deduplicate output operands, and return the
   // - Mapping from operand position in the original op, to operand position in
   // the canonicalized op.
   // - The preserved output operands list (by reference).
   llvm::SmallDenseMap<unsigned, unsigned>
-  deduplicateOutputOperands(GenericOp genericOp,
+  deduplicateOutputOperands(MLIRContext *context, GenericOp genericOp,
                             SmallVector<OpOperand *> &droppedOpOperands,
                             SmallVector<Value> &newOutputOperands,
                             SmallVector<AffineMap> &newIndexingMaps) const {
     llvm::SmallDenseMap<unsigned, unsigned> origToNewPos;
     llvm::SmallDenseMap<std::tuple<Value, AffineMap, Value>, unsigned>
         dedupedOutpts;
     // If the op doesnt have tensor semantics, keep all the outputs as
     // preserved.
for (const auto &outputOpOperand :
       auto key = std::make_tuple(outputOpOperand.value()->get(), indexingMap,
                                  yieldOp->getOperand(outputOpOperand.index()));
       if (isResultValueDead(genericOp, result)) {
         // Check if the opoperand can be dropped without affecting loop
         // bound computation. Add the operand to the list of dropped op
         // operand for checking. If it cannot be dropped, need to pop the
         // value back.
         droppedOpOperands.push_back(outputOpOperand.value());
-        if (genericOp.canOpOperandsBeDropped(droppedOpOperands)) {
+        if (genericOp.canOpOperandsBeDropped(context, droppedOpOperands)) {
           continue;
         }
         droppedOpOperands.pop_back();
       }
 
       if (!genericOp.payloadUsesValueFromOperand(outputOpOperand.value())) {
         // The out operand can also be dropped if it is computed redundantly
         // by another result, the conditions for that are

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

FailureOr<SplitReductionResult> mlir::linalg::splitReduction(
   unsigned insertSplitDimension = control.index;
   if (ratio <= 1)
     return b.notifyMatchFailure(op, "split ratio needs to be greater than 1");
 
   SmallVector<unsigned> dims;
   op.getReductionDims(dims);
   assert(dims.size() == 1);
   unsigned reductionDim = dims[0];
-  SmallVector<int64_t, 4> loopRanges = op.getStaticLoopRanges();
+  SmallVector<int64_t, 4> loopRanges = op.getStaticLoopRanges(b.getContext());
   int64_t reductionDimSize = loopRanges[reductionDim];
   if (reductionDimSize == ShapedType::kDynamicSize ||
       reductionDimSize % ratio != 0 ||
       insertSplitDimension >= loopRanges.size())
     return b.notifyMatchFailure(
         op, "Reduction dimension not divisible by split ratio");
 
   SmallVector<Operation *, 4> combinerOps;
if (splitFactor <= 1)
     return b.notifyMatchFailure(op, "split factor needs to be greater than 1");
 
   SmallVector<unsigned> dims;
   op.getReductionDims(dims);
   if (dims.empty())
     return b.notifyMatchFailure(op, "needs at least 1 reduction dimension");
 
   unsigned reductionDimPos = dims[0];
-  SmallVector<int64_t> loopRanges = op.getStaticLoopRanges();
+  SmallVector<int64_t> loopRanges = op.getStaticLoopRanges(b.getContext());
   int64_t reductionDimSize = loopRanges[reductionDimPos];
   if (reductionDimSize == ShapedType::kDynamicSize ||
       reductionDimSize % splitFactor != 0 ||
       insertSplitDimension >= loopRanges.size())
     return b.notifyMatchFailure(
         op, "first reduction dimension not divisible by split factor");
 
   SmallVector<Operation *> combinerOps;

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

mlir::linalg::computeMultiTileSizes(OpBuilder &builder, LinalgOp op,
   Value targetSizeValue =
       getValueOrCreateConstantIndexOp(builder, loc, targetSize);
   Value divisorValue = getValueOrCreateConstantIndexOp(builder, loc, divisor);
 
   // Find the trip count of the iteration space dimension for which the tile
   // sizes are computed.
   SmallVector<OpFoldResult> allShapes =
       op.createFlatListOfOperandDims(b, b.getLoc());
-  AffineMap shapesToLoops = op.getShapesToLoopsMap();
+  AffineMap shapesToLoops = op.getShapesToLoopsMap(b.getContext());
   SmallVector<OpFoldResult> loopRanges =
       makeComposedFoldedMultiResultAffineApply(b, op.getLoc(), shapesToLoops,
                                                allShapes);
   Value tripCount =
       getValueOrCreateConstantIndexOp(b, op.getLoc(), loopRanges[dimension]);
 
   // Compute the tile sizes and the respective numbers of tiles.
   AffineExpr s0 = b.getAffineSymbolExpr(0);
if (llvm::all_of(tileSizes, isZero)) {
     tiledOp.tensorResults.assign(tiledOp.op->result_begin(),
                                  tiledOp.op->result_end());
     return tiledOp;
   }
 
   // 1. Build the tiled loop ranges.
   SmallVector<OpFoldResult> allShapeSizes =
       op.createFlatListOfOperandDims(b, op.getLoc());
-  AffineMap shapeSizesToLoopsMap = op.getShapesToLoopsMap();
+  AffineMap shapeSizesToLoopsMap = op.getShapesToLoopsMap(b.getContext());
   if (!shapeSizesToLoopsMap)
     return failure();
 
   auto [loopRanges, loopIndexToRangeIndex] = makeTiledLoopRanges(
       b, op.getLoc(), shapeSizesToLoopsMap, allShapeSizes, tileSizes);
 
   SmallVector<utils::IteratorType, 4> iteratorTypes;
   for (const auto &attr : enumerate(op.getIteratorTypesArray())) {

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

struct LinalgOpTilingInterface
   /// Return the iteration domain range.
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
     Location loc = op->getLoc();
     LinalgOp linalgOp = cast<LinalgOp>(op);
     SmallVector<OpFoldResult> allShapesSizes =
         linalgOp.createFlatListOfOperandDims(b, loc);
-    AffineMap map = linalgOp.getShapesToLoopsMap();
+    AffineMap map = linalgOp.getShapesToLoopsMap(b.getContext());
 
     return llvm::to_vector(
         llvm::map_range(map.getResults(), [&](AffineExpr loopExpr) {
           OpFoldResult ofr =
               makeComposedFoldedAffineApply(b, loc, loopExpr, allShapesSizes);
           return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)};
         }));
   }

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

 /// CustomVectorizationHook.
 static VectorizationResult vectorizeLinalgIndex(OpBuilder &b, Operation *op,
                                                 LinalgOp linalgOp) {
   IndexOp indexOp = dyn_cast<linalg::IndexOp>(op);
   if (!indexOp)
     return VectorizationResult{VectorizationStatus::Failure, nullptr};
   auto loc = indexOp.getLoc();
   // Compute the static loop sizes of the index op.
-  auto targetShape = linalgOp.computeStaticLoopSizes();
+  auto targetShape = linalgOp.computeStaticLoopSizes(b.getContext());
   // Compute a one-dimensional index vector for the index op dimension.
   SmallVector<int64_t> constantSeq =
       llvm::to_vector<16>(llvm::seq<int64_t>(0, targetShape[indexOp.getDim()]));
   auto constantOp =
       b.create<arith::ConstantOp>(loc, b.getIndexVectorAttr(constantSeq));
   // Return the one-dimensional index vector if it lives in the trailing
   // dimension of the iteration space since the vectorization algorithm in this
   // case can handle the broadcast.
vectorizeTensorExtract(OpBuilder &b, Operation *op, LinalgOp linalgOp,
   auto loc = extractOp.getLoc();
 
   // Currently only supports extraction with an 1-D index. Checked in the
   // tensorExtractVectorizationPrecondition.
   assert(extractOp.getIndices().size() == 1);
 
   auto indexVec = bvm.lookup(extractOp.getIndices()[0]);
   // Compute the static loop sizes of the extract op.
-  auto targetShape = linalgOp.computeStaticLoopSizes();
+  auto targetShape = linalgOp.computeStaticLoopSizes(b.getContext());
 
   SmallVector<Value> gatherIndices;
   gatherIndices.push_back(b.create<arith::ConstantIndexOp>(loc, 0));
 
   auto maskConstantOp = b.create<arith::ConstantOp>(
       loc,
       DenseIntElementsAttr::get(VectorType::get(targetShape, b.getI1Type()),
                                 /*value=*/true));
vectorizeAsLinalgGeneric(OpBuilder &b, LinalgOp linalgOp,
   bvm.map(valuesSet.getArrayRef(), valuesSet.getArrayRef());
 
   if (linalgOp.getNumDpsInits() == 0)
     return failure();
 
   // TODO: the common vector shape is equal to the static loop sizes only when
   // all indexing maps are projected permutations. For convs and stencils the
   // logic will need to evolve.
-  SmallVector<int64_t> commonVectorShape = linalgOp.computeStaticLoopSizes();
+  SmallVector<int64_t> commonVectorShape =
+      linalgOp.computeStaticLoopSizes(b.getContext());
 
   // 3. Turn all BBArgs into vector.transfer_read / load.
   Location loc = linalgOp.getLoc();
   Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
   for (OpOperand *opOperand : linalgOp.getOpOperandsMatchingBBargs()) {
     BlockArgument bbarg = linalgOp.getMatchingBlockArgument(opOperand);
     if (linalgOp.isScalar(opOperand)) {
       bvm.map(bbarg, opOperand->get());

mlir/lib/IR/AffineMap.cpp

for (unsigned i : llvm::seq(unsigned(0), map.getNumResults())) {
     }
 
     // Reverse each dimension existing in the original map result.
     exprs[map.getDimPosition(i)] = getAffineDimExpr(i, context);
   }
   return AffineMap::get(map.getNumResults(), /*symbolCount=*/0, exprs, context);
 }
 
-AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
+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();
   SmallVector<AffineExpr, 8> results;
   results.reserve(numResults);
   for (auto m : maps) {
     for (auto res : m.getResults())
       results.push_back(res.shiftSymbols(m.getNumSymbols(), numSymbols));
 
     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);
 }
 
 AffineMap mlir::getProjectedMap(AffineMap map,
                                 const llvm::SmallBitVector &unusedDims) {
   return compressUnusedSymbols(compressDims(map, unusedDims));
 }
 
 llvm::SmallBitVector mlir::getUnusedDimsBitVector(ArrayRef<AffineMap> maps) {