[mlir][Linalg] Change signature of `get(Parallel/Reduce/Window)Dims` method.

These method currently takes a SmallVector<AffineExpr> & as an
argument to return the dims as AffineExpr. This creation of
AffineExpr objects is unnecessary.

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

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

@@ -90,7 +90,7 @@ SmallVector<AffineExpr, 4> concat(ArrayRef<AffineExpr> a,
 /// Return the dims that are `iteratorTypeName` loops in the LinalgOp `op`.
 /// Assumes `op` is a LinalgOp.
 void getDimsOfType(Operation *op, StringRef iteratorTypeName,
-                   SmallVectorImpl<AffineExpr> &res);
+                   SmallVectorImpl<unsigned> &res);
 
 namespace detail {
 LogicalResult verifyStructuredOpInterface(Operation *op);

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

@@ -158,7 +158,7 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
       }],
       /*retTy=*/"void",
       /*methodName=*/"getParallelDims",
-      /*args=*/(ins "SmallVectorImpl<AffineExpr> &":$res),
+      /*args=*/(ins "SmallVectorImpl<unsigned> &":$res),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return getDimsOfType($_op, getParallelIteratorTypeName(), res);
@@ -183,7 +183,7 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
       }],
       /*retTy=*/"void",
       /*methodName=*/"getReductionDims",
-      /*args=*/(ins "SmallVectorImpl<AffineExpr> &":$res),
+      /*args=*/(ins "SmallVectorImpl<unsigned> &":$res),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return getDimsOfType($_op, getReductionIteratorTypeName(), res);
@@ -208,7 +208,7 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
       }],
       /*retTy=*/"void",
       /*methodName=*/"getWindowDims",
-      /*args=*/(ins "SmallVectorImpl<AffineExpr> &":$res),
+      /*args=*/(ins "SmallVectorImpl<unsigned> &":$res),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return getDimsOfType($_op.getOperation(), getWindowIteratorTypeName(), res);

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

@@ -612,7 +612,7 @@ LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {
              << indexingMap.getNumResults() << ")";
   }
 
-  SmallVector<AffineExpr> redDims;
+  SmallVector<unsigned> redDims;
   linalgOp.getReductionDims(redDims);
 
   // Simplifying assumption: either full tensor or full buffer mode.
@@ -638,9 +638,8 @@ LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {
   // Output tensor indexing map may not depend on reduction indices.
   for (OpOperand *opOperand : linalgOp.getOutputOperands()) {
     AffineMap indexingMap = linalgOp.getTiedIndexingMap(opOperand);
-    for (auto expr : indexingMap.getResults()) {
-      for (auto dim : redDims) {
-        unsigned pos = dim.cast<AffineDimExpr>().getPosition();
+    for (AffineExpr expr : indexingMap.getResults()) {
+      for (unsigned pos : redDims) {
         if (expr.isFunctionOfDim(pos)) {
           std::string exprStr;
           {

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

@@ -2318,16 +2318,15 @@ static LogicalResult verify(IndexOp op) {
 /// Return the dims that are `iteratorTypeName` loops in the LinalgOp `op`.
 /// Assumes `op` is a LinalgOp.
 void mlir::linalg::getDimsOfType(Operation *op, StringRef iteratorTypeName,
-                                 SmallVectorImpl<AffineExpr> &res) {
+                                 SmallVectorImpl<unsigned> &res) {
   if (!cast<LinalgOp>(op).iterator_types())
     return;
 
   unsigned dim = 0;
-  MLIRContext *ctx = op->getContext();
   for (auto tn :
        cast<LinalgOp>(op).iterator_types().getAsValueRange<StringAttr>()) {
     if (tn == iteratorTypeName)
-      res.push_back(getAffineDimExpr(dim, ctx));
+      res.push_back(dim);
     ++dim;
   }
 }