[mlir][linalg] Cleanup LinalgOp usage in drop unit dims.

[gysit]

Summary

Replace the uses of deprecated Structured Op Interface methods in DropUnitDims.cpp. This patch is based on https://reviews.llvm.org/D103394.

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

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));
     if (!invertedMap)
       return failure();
-    SmallVector<int64_t, 4> dims;
-    for (ShapedType shapedType : genericOp.getShapedOperandTypes())
-      dims.append(shapedType.getShape().begin(), shapedType.getShape().end());
+    SmallVector<int64_t> dims = genericOp.getStaticShape();
 
     // Find all the reduction iterators. Those need some special consideration
     // (see below).
     auto getLoopDimsOfType =
         [&](StringRef iteratorTypeName) -> SmallVector<unsigned, 4> {
       SmallVector<AffineExpr> dimExprs;
       getDimsOfType(genericOp, iteratorTypeName, dimExprs);
       return llvm::to_vector<4>(llvm::map_range(dimExprs, [](AffineExpr expr) {

 /// an operand/result with the unit-extent dimension removed. This is only done
 /// if the indexing map used to access that didimensionmension has a
 /// AffineConstantExpr of value 0. Given the `type` of an result/operand of a
 /// Linalg op, and its `indexMap` the utility function returns:
 /// - the new type with dimensions of size 1 removed.
 /// - modified index map that can be used to access the replaced result/operand
 /// - the reassociation that converts from the original tensor type to the
 ///   modified tensor type.
-static UnitExtentReplacementInfo replaceUnitExtents(AffineMap indexMap,
-                                                    RankedTensorType type,
+static UnitExtentReplacementInfo replaceUnitExtents(GenericOp genericOp,
+                                                    OpOperand *opOperand,
                                                     MLIRContext *context) {
-  ArrayRef<int64_t> shape = type.getShape();
-  ArrayRef<AffineExpr> exprs = indexMap.getResults();
+  AffineMap indexingMap = genericOp.getTiedIndexingMap(opOperand);
+  ArrayRef<int64_t> shape = genericOp.getShape(opOperand);
+  ArrayRef<AffineExpr> exprs = indexingMap.getResults();
   SmallVector<AffineExpr, 2> reassociations;
   SmallVector<Attribute, 4> reassociationMaps;
   SmallVector<AffineExpr, 4> newIndexExprs;
   SmallVector<int64_t, 4> newShape;
 
-  int64_t origRank = type.getRank();
+  int64_t origRank = genericOp.getRank(opOperand);
   AffineExpr zeroExpr = getAffineConstantExpr(0, context);
   auto isUnitExtent = [&](int64_t dim) -> bool {
     return shape[dim] == 1 && exprs[dim] == zeroExpr;
   };
 
   unsigned dim = 0;
   // Fold dimensions that are unit-extent at the beginning of the tensor.
   while (dim < origRank && isUnitExtent(dim))
     reassociations.push_back(getAffineDimExpr(dim++, context));
   while (dim < origRank) {
     reassociations.push_back(getAffineDimExpr(dim, context));
     newIndexExprs.push_back(exprs[dim]);
     newShape.push_back(shape[dim]);
     // Fold all following dimensions that are unit-extent.
     while (dim + 1 < origRank && isUnitExtent(dim + 1)) {
       ++dim;
       reassociations.push_back(getAffineDimExpr(dim, context));
     }
     reassociationMaps.push_back(AffineMapAttr::get(AffineMap::get(
         origRank, /*symbolCount = */ 0, reassociations, context)));
     reassociations.clear();
     ++dim;
   }
   UnitExtentReplacementInfo info = {
-      RankedTensorType::get(newShape, type.getElementType()),
-      AffineMap::get(indexMap.getNumDims(), indexMap.getNumSymbols(),
+      RankedTensorType::get(newShape,
+                            getElementTypeOrSelf(opOperand->get().getType())),
+      AffineMap::get(indexingMap.getNumDims(), indexingMap.getNumSymbols(),
                      newIndexExprs, context),
       ArrayAttr::get(context, reassociationMaps)};
   return info;
 }
 
 namespace {
 
 SmallVector<ReassociationExprs, 2>
LogicalResult matchAndRewrite(GenericOp genericOp,
 
     MLIRContext *context = rewriter.getContext();
     Location loc = genericOp.getLoc();
 
     SmallVector<AffineMap, 4> newIndexingMaps;
     SmallVector<ArrayAttr, 4> reassociationMaps;
     SmallVector<ShapedType, 4> newInputOutputTypes;
     bool doCanonicalization = false;
-    for (auto it : llvm::zip(genericOp.getIndexingMaps(),
-                             genericOp.getShapedOperandTypes())) {
-      auto replacementInfo = replaceUnitExtents(
-          std::get<0>(it), std::get<1>(it).template cast<RankedTensorType>(),
-          context);
+
+    for (OpOperand *opOperand : genericOp.getInputAndOutputOperands()) {
+      auto replacementInfo = replaceUnitExtents(genericOp, opOperand, context);
       reassociationMaps.push_back(replacementInfo.reassociation);
       newIndexingMaps.push_back(replacementInfo.indexMap);
       newInputOutputTypes.push_back(replacementInfo.type);
-      doCanonicalization |= replacementInfo.type != std::get<1>(it);
+      doCanonicalization |= replacementInfo.type != opOperand->get().getType();
     }
 
     // If the indexing maps of the result operation are not invertible (i.e. not
     // legal), abort.
     if (!doCanonicalization ||
         !inversePermutation(concatAffineMaps(newIndexingMaps)))
       return failure();