[mlir] Move getDimsOfType to StructuredOpsUtils.h.

[olegshyshkov]

Summary

This change will bring all helpers that work with iterator types to one place.
Currently getDimsOfType is is declared in Linalg.h, but not directly included by
LinalgInterfaces. It worked so far only because all the places that include
LinalgInterfaces.h also include Linalg.h directly or indirectly.

[olegshyshkov]

Fix typo.

[pifon2a]

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

261–262

If you're not passing the operation itself, then the function becomes more general and just computes positions of x: type in vector<type>. In that case, you probably want to rename it to findPositions or smth like that.

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

117

nit: use (auto en : llvm::enumerate(iteratorType...)) and then en.index() and en.value()

[olegshyshkov]

Resolved comments.

[olegshyshkov]

Fix indent.

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

 /// symbol-less identity map of `rank`.
 AffineMap extractOrIdentityMap(Optional<AffineMap> maybeMap, unsigned rank,
                                MLIRContext *context);
 
 /// Return the vector that is the concatenation of `a` and `b`.
 SmallVector<AffineExpr, 4> concat(ArrayRef<AffineExpr> a,
                                   ArrayRef<AffineExpr> b);
 
-/// Return the dims that are `iteratorTypeName` loops in the LinalgOp `op`.
-/// Assumes `op` is a LinalgOp.
-void getDimsOfType(Operation *op, StringRef iteratorTypeName,
-                   SmallVectorImpl<unsigned> &res);
-
 } // namespace linalg
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Linalg Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgOpsDialect.h.inc"

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

InterfaceMethod<
       /*desc=*/[{
         Return the dims that are parallel loops.
       }],
       /*retTy=*/"void",
       /*methodName=*/"getParallelDims",
       /*args=*/(ins "SmallVectorImpl<unsigned> &":$res),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return getDimsOfType($_op, getParallelIteratorTypeName(), res);
+        return findPositionsOfType($_op.iterator_types(),
+                                   getParallelIteratorTypeName(), res);
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the number of reduction loops.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getNumReductionLoops",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return getNumIterators(getReductionIteratorTypeName(),
                                $_op.iterator_types());
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the dims that are reduction loops.
       }],
       /*retTy=*/"void",
       /*methodName=*/"getReductionDims",
       /*args=*/(ins "SmallVectorImpl<unsigned> &":$res),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return getDimsOfType($_op, getReductionIteratorTypeName(), res);
+        return findPositionsOfType($_op.iterator_types(),
+                                   getReductionIteratorTypeName(), res);
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the number of window loops.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getNumWindowLoops",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return getNumIterators(getWindowIteratorTypeName(),
                                $_op.iterator_types());
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the dims that are window loops.
       }],
       /*retTy=*/"void",
       /*methodName=*/"getWindowDims",
       /*args=*/(ins "SmallVectorImpl<unsigned> &":$res),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return getDimsOfType($_op.getOperation(), getWindowIteratorTypeName(), res);
+        return findPositionsOfType($_op.iterator_types(),
+                                   getWindowIteratorTypeName(), res);
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the total number of loops within the current operation.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getNumLoops",

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

 
 inline unsigned getNumIterators(ArrayAttr iteratorTypes) {
   unsigned res = 0;
   for (auto n : getAllIteratorTypeNames())
     res += getNumIterators(n, iteratorTypes);
   return res;
 }
 
+/// Return positions in `iteratorTypes` that match `iteratorTypeName`.
+inline void findPositionsOfType(ArrayAttr iteratorTypes,
+                                StringRef iteratorTypeName,
+                                SmallVectorImpl<unsigned> &res) {
+  for (const auto &en :
+       llvm::enumerate(iteratorTypes.getAsValueRange<StringAttr>())) {
+    if (en.value() == iteratorTypeName)
+      res.push_back(en.index());
+  }
+}
+
 /// Helper StructuredGenerator class to manipulate and rewrite ops with
 /// `StructuredOpInterface`. This is templated for now because VectorOps do not
 /// yet implement the StructuredOpInterface itself.
 template <typename StructuredOpInterface>
 class StructuredGenerator {
 public:
   using MapList = ArrayRef<ArrayRef<AffineExpr>>;
 

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

 #include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
 
-/// 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<unsigned> &res) {
-  if (!cast<LinalgOp>(op).iterator_types())
-    return;
-
-  unsigned dim = 0;
-  for (auto tn :
-       cast<LinalgOp>(op).iterator_types().getAsValueRange<StringAttr>()) {
-    if (tn == iteratorTypeName)
-      res.push_back(dim);
-    ++dim;
-  }
-}
-
 AffineMap mlir::linalg::extractOrIdentityMap(Optional<AffineMap> maybeMap,
                                              unsigned rank,
                                              MLIRContext *context) {
   if (maybeMap)
     return *maybeMap;
   if (rank == 0)
     return AffineMap::get(context);
   return AffineMap::getMultiDimIdentityMap(rank, context);