[mlir] Move helper function to work with IteratorTypes to StructuredOpsUtils.

[olegshyshkov]

Summary

isParallelIterator and isReductionIterator are moved from Linalg/Utils to
the new location so they can be used in LinalgOps itself.

getNParallelIterators is a very common pattern that was recreated 3 times in
the MLIR codebase.

[mravishankar]

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

125

I find it strange that there is a utils namespace for these, but other methods here are not in utils namespace. Either we should move all of these to utils namespace, or drop the namespace altogether.

139

I am not sure this is anything more than return SmallVector<StringRef>(nParallel, getParallelIteratorType(). Dont think that needs a utility function.

[olegshyshkov]

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

125

I totally agree that it's not great and everything in this file should be under under. I checked what other methods are in this file and looks like there is more clean up to be done.

• Getters for attribute names should probably go away, because they duplicate what is already generated by tablegen.
• Getters for IteratorType names will got away once I'm done with this RFC.

So there is no reason to go over the same places in code twice. Would you unblock this patch if I leave a TODO and clean up in the following changes?

139

It is indeed just a SmallVector construction, but in my opinion it greatly improves readability and hides unnecessary verbose details when ops are constructed. This function was copied multiple time to different parts of the codebase and is heavily used, so I would think it's useful.

Would you suggest to inline all the existing usages instead?

[mravishankar]

Getters for attribute names should probably go away, because they duplicate what is already generated by tablegen.

+1!

Getters for IteratorType names will got away once I'm done with this RFC.

+1!

Would you suggest to inline all the existing usages instead?

That would be great. Lesser API surface area, the better.

Would you unblock this patch if I leave a TODO and clean up in the following changes?

I am only looking for consistency. In other dialects that provide utils (like Arith dialect and Tensor dialect) I dont see the util namespace being used (is util used anywhere in the codebase?). It just makes it confusing for developers looking around w.r.t. when to add util and when not to do so.
I'd suggest you just drop the util and save yourself some headache (unless you really care about it and want to establish some guidelines across MLIR about using of util namespace, which would also be fine by me)

[olegshyshkov]

I expected that this patch will help reduce the size of https://reviews.llvm.org/D137658, but the gain is so small, I don't think it's worth proceeding. I'll just drop this change.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 bool allIndexingsAreProjectedPermutation(LinalgOp op);
 
 /// Detect whether `r` has only ConstantOp, ElementwiseMappable and YieldOp.
 bool hasOnlyScalarElementwiseOp(Region &r);
 
 /// Check if a LinalgOp is an element-wise operation.
 bool isElementwise(LinalgOp op);
 
-/// Check if iterator type has "parallel" semantics.
-bool isParallelIterator(StringRef iteratorType);
-
-/// Check if iterator type  has "reduction" semantics.
-bool isReductionIterator(StringRef iteratorType);
-
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
 OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
                                int64_t dim);
 
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.

mlir/include/mlir/Dialect/Tosa/Utils/CoversionUtils.h

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 namespace tosa {
 
-// Creates a SmallVector of Stringrefs for N parallel loops
-SmallVector<StringRef> getNParallelLoopsAttrs(unsigned nParallelLoops);
-
 // Takes a vector of values and condenses them to a vector with no gaps.
 SmallVector<Value> condenseValues(const SmallVector<Value> &values);
 
 // Takes the parameters for a clamp and turns it into a series of ops for float
 // inputs.
 Value clampFloatHelper(Location loc, Value arg, arith::ConstantOp min,
                        arith::ConstantOp max, OpBuilder &rewriter);
 

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

inline void findPositionsOfType(ArrayRef<StringRef> iteratorTypes,
                                 StringRef iteratorTypeName,
                                 SmallVectorImpl<unsigned> &res) {
   for (const auto &en : llvm::enumerate(iteratorTypes)) {
     if (en.value() == iteratorTypeName)
       res.push_back(en.index());
   }
 }
 
+namespace utils {
+
+/// Check if iterator type has "parallel" semantics.
+inline bool isParallelIterator(StringRef iteratorType) {
+  return iteratorType == getParallelIteratorTypeName();
+}
+
+/// Check if iterator type  has "reduction" semantics.
+inline bool isReductionIterator(StringRef iteratorType) {
+  return iteratorType == getReductionIteratorTypeName();
+}
+
+/// Returns an array of iterator types that contains `nParallelLoops` "parallel"
+/// attributes.
+SmallVector<StringRef> getNParallelIterators(unsigned nParallel);
+
+} // namespace utils
+
 /// 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/Conversion/TosaToLinalg/TosaToLinalg.cpp

elementwiseMatchAndRewriteHelper(Operation *operation,
   }
 
   indexingMaps.append(operation->getNumResults(),
                       rewriter.getMultiDimIdentityMap(rank));
 
   bool didEncounterError = false;
   auto linalgOp = rewriter.create<linalg::GenericOp>(
       loc, opResultTypes, operands, emptyTensors, indexingMaps,
-      getNParallelLoopsAttrs(rank),
+      utils::getNParallelIterators(rank),
       [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) {
         Value opResult = createLinalgBodyCalculationForElementwiseOp(
             operation, blockArgs.take_front(operation->getNumOperands()),
             bodyResultTypes, rewriter);
         if (!opResult) {
           didEncounterError = true;
           return;
         }
LogicalResult matchAndRewrite(tosa::TransposeOp op,
 
     SmallVector<AffineMap, 2> affineMaps = {
         AffineMap::get(resultTy.getRank(), /*symbolCount=*/0, inputExprs,
                        rewriter.getContext()),
         rewriter.getMultiDimIdentityMap(resultTy.getRank())};
 
     rewriter.replaceOpWithNewOp<linalg::GenericOp>(
         op, resultTy, op.getInput1(), ValueRange{emptyTensor}, affineMaps,
-        getNParallelLoopsAttrs(resultTy.getRank()),
+        utils::getNParallelIterators(resultTy.getRank()),
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
           nestedBuilder.create<linalg::YieldOp>(loc, *args.begin());
         });
     return success();
   }
 };
 
 class RescaleConverter : public OpRewritePattern<tosa::RescaleOp> {
LogicalResult matchAndRewrite(tosa::RescaleOp op,
 
     // Construct the indexing maps needed for linalg.generic ops.
     Value emptyTensor = rewriter.create<tensor::EmptyOp>(
         loc, outputTy.getShape(), outputTy.getElementType(),
         ArrayRef<Value>({dynDims}));
 
     auto linalgOp = rewriter.create<linalg::GenericOp>(
         loc, outputTy, genericInputs, ValueRange{emptyTensor}, indexingMaps,
-        getNParallelLoopsAttrs(rank),
+        utils::getNParallelIterators(rank),
         [&](OpBuilder &nestedBuilder, Location nestedLoc,
             ValueRange blockArgs) {
           Value value = blockArgs[0];
           Type valueTy = value.getType();
 
           // For now we do all of our math in 64-bit. This is not optimal but
           // should be correct for now, consider computing correct bit depth
           // later.
auto emptyTensor = rewriter.create<tensor::EmptyOp>(
         loc, resultTy.getShape(), resultElementTy, dynamicDims);
 
     SmallVector<AffineMap, 2> affineMaps = {
         rewriter.getMultiDimIdentityMap(resultTy.getRank())};
 
     Value resize = input;
     auto genericOp = rewriter.create<linalg::GenericOp>(
         loc, resultTy, ValueRange({}), ValueRange{emptyTensor}, affineMaps,
-        getNParallelLoopsAttrs(resultTy.getRank()));
+        utils::getNParallelIterators(resultTy.getRank()));
     resize = genericOp.getResult(0);
 
     OpBuilder::InsertionGuard regionGuard(rewriter);
     rewriter.createBlock(&genericOp.getRegion(), genericOp.getRegion().end(),
                          TypeRange({resultElementTy}), loc);
     Value batch = rewriter.create<linalg::IndexOp>(loc, 0);
     Value y = rewriter.create<linalg::IndexOp>(loc, 1);
     Value x = rewriter.create<linalg::IndexOp>(loc, 2);
auto emptyTensor = rewriter
                                                     inputTy.getElementType(),
                                                     ArrayRef<Value>({dynDims}))
                            .getResult();
     SmallVector<AffineMap, 2> affineMaps = {
         rewriter.getMultiDimIdentityMap(resultTy.getRank())};
 
     rewriter.replaceOpWithNewOp<linalg::GenericOp>(
         op, resultTy, ArrayRef<Value>({}), ValueRange{emptyTensor}, affineMaps,
-        getNParallelLoopsAttrs(resultTy.getRank()),
+        utils::getNParallelIterators(resultTy.getRank()),
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
           llvm::SmallVector<Value> indices;
           for (unsigned int i = 0; i < inputTy.getRank(); i++) {
             auto index =
                 rewriter.create<linalg::IndexOp>(nestedLoc, i).getResult();
             if (i == axis) {
               auto one = rewriter.create<arith::ConstantIndexOp>(nestedLoc, 1);
               auto sizeMinusOne =
auto readAffineMap =
                        rewriter.getContext());
 
     SmallVector<AffineMap, 2> affineMaps = {
         readAffineMap, rewriter.getMultiDimIdentityMap(genericShape.size())};
 
     auto genericOp = rewriter.create<linalg::GenericOp>(
         loc, RankedTensorType::get(genericShape, elementTy), input,
         ValueRange{emptyTensor}, affineMaps,
-        getNParallelLoopsAttrs(genericShape.size()),
+        utils::getNParallelIterators(genericShape.size()),
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
           nestedBuilder.create<linalg::YieldOp>(op.getLoc(), *args.begin());
         });
 
     rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
         op, resultTy, genericOp.getResult(0),
         rewriter.getI64ArrayAttr(resultTy.getShape()));
     return success();
LogicalResult matchAndRewrite(tosa::ArgMaxOp argmaxOp,
     auto filledTensorMax =
         rewriter
             .create<linalg::FillOp>(loc, ValueRange{fillValueMax},
                                     ValueRange{emptyTensorMax})
             .result();
 
     // We need to reduce along the arg-max axis, with parallel operations along
     // the rest.
-    SmallVector<StringRef, 4> iteratorTypes;
-    iteratorTypes.resize(inputTy.getRank(), getParallelIteratorTypeName());
+    SmallVector<StringRef, 4> iteratorTypes =
+        utils::getNParallelIterators(inputTy.getRank());
     iteratorTypes[axis] = getReductionIteratorTypeName();
 
     SmallVector<AffineExpr, 2> srcExprs;
     SmallVector<AffineExpr, 2> dstExprs;
     for (int i = 0, rank = inputTy.getRank(); i != rank; ++i) {
       srcExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext()));
       if (axis != i)
         dstExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext()));
SmallVector<AffineMap, 2> affineMaps = {
             /*dimCount=*/resultTy.getRank(), /*symbolCount=*/0,
             {rewriter.getAffineDimExpr(0), rewriter.getAffineDimExpr(1)},
             rewriter.getContext()),
         rewriter.getMultiDimIdentityMap(resultTy.getRank())};
 
     auto genericOp = rewriter.create<linalg::GenericOp>(
         loc, ArrayRef<Type>({resultTy}), ValueRange{indices},
         ValueRange{emptyTensor}, affineMaps,
-        getNParallelLoopsAttrs(resultTy.getRank()),
+        utils::getNParallelIterators(resultTy.getRank()),
         [&](OpBuilder &b, Location loc, ValueRange args) {
           auto indexValue = args[0];
           auto index0 = rewriter.create<linalg::IndexOp>(loc, 0);
           Value index1 = rewriter.create<arith::IndexCastOp>(
               loc, rewriter.getIndexType(), indexValue);
           auto index2 = rewriter.create<linalg::IndexOp>(loc, 2);
           Value extract = rewriter.create<tensor::ExtractOp>(
               loc, input, ValueRange{index0, index1, index2});
auto emptyTensor = rewriter
                            .getResult();
 
     SmallVector<AffineMap, 2> affineMaps = {
         rewriter.getMultiDimIdentityMap(resultTy.getRank()),
         rewriter.getMultiDimIdentityMap(resultTy.getRank())};
 
     auto genericOp = rewriter.create<linalg::GenericOp>(
         loc, resultTy, ValueRange({input}), ValueRange{emptyTensor}, affineMaps,
-        getNParallelLoopsAttrs(resultTy.getRank()));
+        utils::getNParallelIterators(resultTy.getRank()));
     rewriter.replaceOp(op, genericOp.getResult(0));
 
     {
       OpBuilder::InsertionGuard regionGuard(rewriter);
       Block *block = rewriter.createBlock(
           &genericOp.getRegion(), genericOp.getRegion().end(),
           TypeRange({inputElementTy, resultElementTy}), {loc, loc});
 

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

if (isQuantized) {
                   loc, resultTy, ValueRange{input, weight, iZpVal, kZpVal},
                   ValueRange{zeroTensor}, strideAttr, dilationAttr)
               ->getResult(0);
 
       Value result =
           rewriter
               .create<linalg::GenericOp>(
                   loc, resultTy, ValueRange({bias, conv}), biasEmptyTensor,
-                  indexingMaps, getNParallelLoopsAttrs(resultTy.getRank()),
+                  indexingMaps,
+                  utils::getNParallelIterators(resultTy.getRank()),
                   [&](OpBuilder &nestedBuilder, Location nestedLoc,
                       ValueRange args) {
                     Value added = nestedBuilder.create<arith::AddIOp>(
                         loc, args[0], args[1]);
                     nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
                   })
               .getResult(0);
       rewriter.replaceOp(op, result);
       return success();
     }
 
     Value conv = rewriter
                      .create<linalg::Conv2DNhwcHwcfOp>(
                          loc, resultTy, ValueRange{input, weight},
                          ValueRange{zeroTensor}, strideAttr, dilationAttr)
                      ->getResult(0);
 
     Value result =
         rewriter
             .create<linalg::GenericOp>(
                 loc, resultTy, ValueRange({bias, conv}), biasEmptyTensor,
-                indexingMaps, getNParallelLoopsAttrs(resultTy.getRank()),
+                indexingMaps, utils::getNParallelIterators(resultTy.getRank()),
                 [&](OpBuilder &nestedBuilder, Location nestedLoc,
                     ValueRange args) {
                   Value added = nestedBuilder.create<arith::AddFOp>(
                       loc, args[0], args[1]);
                   nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
                 })
             .getResult(0);
 
if (!isQuantized) {
       Value convReshape = rewriter.create<tensor::CollapseShapeOp>(
           loc, resultTy, conv, reassociationMap);
 
       Value result =
           rewriter
               .create<linalg::GenericOp>(
                   loc, resultTy, ValueRange({bias, convReshape}),
                   biasEmptyTensor, indexingMaps,
-                  getNParallelLoopsAttrs(resultRank),
+                  utils::getNParallelIterators(resultRank),
                   [&](OpBuilder &nestedBuilder, Location nestedLoc,
                       ValueRange args) {
                     Value added = nestedBuilder.create<arith::AddFOp>(
                         loc, args[0], args[1]);
                     nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
                   })
               .getResult(0);
       rewriter.replaceOp(op, result);
if (!isQuantized) {
       createDepthwiseConvCollapseMap(resultRank, reassociationMap, rewriter);
       Value convReshape = rewriter.create<tensor::CollapseShapeOp>(
           loc, resultTy, conv, reassociationMap);
       Value result =
           rewriter
               .create<linalg::GenericOp>(
                   loc, resultTy, ValueRange({bias, convReshape}),
                   biasEmptyTensor, indexingMaps,
-                  getNParallelLoopsAttrs(resultRank),
+                  utils::getNParallelIterators(resultRank),
                   [&](OpBuilder &nestedBuilder, Location nestedLoc,
                       ValueRange args) {
                     Value added = nestedBuilder.create<arith::AddIOp>(
                         loc, args[0], args[1]);
                     nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
                   })
               .getResult(0);
       rewriter.replaceOp(op, result);
if (!op.getQuantizationInfo()) {
                              loc, TypeRange{op.getType()},
                              ValueRange{input, transposedWeight}, zeroTensor)
                          ->getResult(0);
 
       Value result =
           rewriter
               .create<linalg::GenericOp>(
                   loc, outputTy, ValueRange({bias, matmul}), biasEmptyTensor,
-                  indexingMaps, getNParallelLoopsAttrs(outputTy.getRank()),
+                  indexingMaps,
+                  utils::getNParallelIterators(outputTy.getRank()),
                   [&](OpBuilder &nestedBuilder, Location nestedLoc,
                       ValueRange args) {
                     Value added = nestedBuilder.create<arith::AddFOp>(
                         loc, args[0], args[1]);
                     nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
                   })
               .getResult(0);
       rewriter.replaceOp(op, result);
Value matmul =
                 loc, TypeRange{op.getType()},
                 ValueRange{input, transposedWeight, inputZp, outputZp},
                 zeroTensor)
             ->getResult(0);
     Value result =
         rewriter
             .create<linalg::GenericOp>(
                 loc, outputTy, ValueRange({bias, matmul}), biasEmptyTensor,
-                indexingMaps, getNParallelLoopsAttrs(outputTy.getRank()),
+                indexingMaps, utils::getNParallelIterators(outputTy.getRank()),
                 [&](OpBuilder &nestedBuilder, Location nestedLoc,
                     ValueRange args) {
                   Value added = nestedBuilder.create<arith::AddIOp>(
                       loc, args[0], args[1]);
                   nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
                 })
             .getResult(0);
     rewriter.replaceOp(op, result);
LogicalResult matchAndRewrite(tosa::AvgPool2dOp op,
 
     Value genericEmptyTensor = rewriter.create<tensor::EmptyOp>(
         loc, resultTy.getShape(), resultETy, dynamicDims);
 
     auto genericOp = rewriter.create<linalg::GenericOp>(
         loc, ArrayRef<Type>({resultTy}), ValueRange{poolingOp},
         ValueRange{genericEmptyTensor},
         ArrayRef<AffineMap>({affineMap, affineMap}),
-        getNParallelLoopsAttrs(resultTy.getRank()),
+        utils::getNParallelIterators(resultTy.getRank()),
         [&](OpBuilder &b, Location loc, ValueRange args) {
           auto zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
           auto one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
           auto iH = rewriter.create<arith::ConstantIndexOp>(
               loc, poolingOpTy.getDimSize(1) - 1);
           auto iW = rewriter.create<arith::ConstantIndexOp>(
               loc, poolingOpTy.getDimSize(2) - 1);
 

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

static MatchConvolutionResult isConvolutionInterfaceImpl(Operation *op) {
   // - Depth multiplier : unconvolved in input, present in output, present in
   //   filter.
   llvm::SmallDenseSet<unsigned> allLoopDims;
   for (auto outputExpr : indexingMaps.back().getResults()) {
     unsigned outputDim = outputExpr.cast<AffineDimExpr>().getPosition();
     if (inputExprWalker.unConvolvedDims.count(outputDim) &&
         !filterDims.count(outputDim)) {
       // Batch dimension.
-      if (iteratorTypes[outputDim] != getParallelIteratorTypeName())
+      if (!utils::isParallelIterator(iteratorTypes[outputDim]))
         return MatchConvolutionResult::OutputDimsNotParallel;
       allLoopDims.insert(outputDim);
       continue;
     }
     if (inputExprWalker.convolvedDims.count(outputDim) &&
         !filterDims.count(outputDim)) {
       // Output image Loop dimension.
-      if (iteratorTypes[outputDim] != getParallelIteratorTypeName())
+      if (!utils::isParallelIterator(iteratorTypes[outputDim]))
         return MatchConvolutionResult::OutputDimsNotParallel;
       allLoopDims.insert(outputDim);
       continue;
     }
     if (!inputExprWalker.convolvedDims.count(outputDim) &&
         !inputExprWalker.unConvolvedDims.count(outputDim) &&
         filterDims.count(outputDim)) {
       // Output channel dimension.
-      if (iteratorTypes[outputDim] != getParallelIteratorTypeName())
+      if (!utils::isParallelIterator(iteratorTypes[outputDim]))
         return MatchConvolutionResult::OutputDimsNotParallel;
       allLoopDims.insert(outputDim);
       continue;
     }
     if (inputExprWalker.unConvolvedDims.count(outputDim) &&
         filterDims.count(outputDim)) {
       // Depth multiplier.
-      if (iteratorTypes[outputDim] != getParallelIteratorTypeName())
+      if (!utils::isParallelIterator(iteratorTypes[outputDim]))
         return MatchConvolutionResult::OutputDimsNotParallel;
       allLoopDims.insert(outputDim);
       continue;
     }
     return MatchConvolutionResult::NonConvolutionLoop;
   }
   for (auto filterExpr : indexingMaps[1].getResults()) {
     unsigned filterDim = filterExpr.cast<AffineDimExpr>().getPosition();
     if (outputDims.count(filterDim) &&
         !inputExprWalker.unConvolvedDims.count(filterDim) &&
         !inputExprWalker.convolvedDims.count(filterDim)) {
       // Output channel dimension. THis is already seen, continue;
       continue;
     }
     if (inputExprWalker.convolvedDims.count(filterDim) &&
         !outputDims.count(filterDim)) {
       // Filter loop dimension.
-      if (iteratorTypes[filterDim] != getReductionIteratorTypeName())
+      if (!utils::isReductionIterator(iteratorTypes[filterDim]))
         return MatchConvolutionResult::NonOutputDimNotReduction;
       if (allLoopDims.count(filterDim))
         return MatchConvolutionResult::NonConvolutionLoop;
       allLoopDims.insert(filterDim);
       continue;
     }
     if (inputExprWalker.unConvolvedDims.count(filterDim) &&
         !outputDims.count(filterDim)) {
       // Input channel dimension.
-      if (iteratorTypes[filterDim] != getReductionIteratorTypeName())
+      if (!utils::isReductionIterator(iteratorTypes[filterDim]))
         return MatchConvolutionResult::NonOutputDimNotReduction;
       if (allLoopDims.count(filterDim))
         return MatchConvolutionResult::NonConvolutionLoop;
       allLoopDims.insert(filterDim);
       continue;
     }
     if (inputExprWalker.unConvolvedDims.count(filterDim) &&
         outputDims.count(filterDim)) {

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

for (Type inputArgType : TypeRange{getInputs()}) {
     }
   }
 
   return success();
 }
 
 SmallVector<StringRef> MapOp::getIteratorTypesArray() {
   int64_t rank = getInit().getType().getRank();
-  return SmallVector<StringRef>(rank, getParallelIteratorTypeName());
+  return utils::getNParallelIterators(rank);
 }
 
 ArrayAttr MapOp::getIndexingMaps() {
   Builder builder(getContext());
   int64_t rank = getInit().getType().getRank();
   int64_t numIndexingMaps = getOperands().size();
   return builder.getAffineMapArrayAttr(SmallVector<AffineMap>(
       numIndexingMaps, builder.getMultiDimIdentityMap(rank)));
void ReduceOp::build(
   }
 
   if (bodyBuild)
     buildGenericRegion(builder, result, inputs, inits, bodyBuild);
 }
 
 SmallVector<StringRef> ReduceOp::getIteratorTypesArray() {
   int64_t inputRank = getInputs()[0].getType().cast<ShapedType>().getRank();
-  SmallVector<StringRef> iteratorTypes(inputRank,
-                                       getParallelIteratorTypeName());
+  SmallVector<StringRef> iteratorTypes =
+      utils::getNParallelIterators(inputRank);
   for (int64_t reductionDim : getDimensions())
     iteratorTypes[reductionDim] = getReductionIteratorTypeName();
   return iteratorTypes;
 }
 
 ArrayAttr ReduceOp::getIndexingMaps() {
   int64_t inputRank = getInputs()[0].getType().cast<ShapedType>().getRank();
   SmallVector<AffineMap> affineMaps(
for (int64_t i = 0; i < rank; ++i) {
     }
   }
 
   return success();
 }
 
 SmallVector<StringRef> TransposeOp::getIteratorTypesArray() {
   int64_t rank = getInit().getType().getRank();
-  return SmallVector<StringRef>(rank, getParallelIteratorTypeName());
+  return utils::getNParallelIterators(rank);
 }
 
 ArrayAttr TransposeOp::getIndexingMaps() {
   Builder builder(getContext());
   int64_t rank = getInit().getType().getRank();
   return builder.getAffineMapArrayAttr(
       {builder.getMultiDimIdentityMap(rank),
        AffineMap::getPermutationMap(

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

static bool isFusableWithReshapeByDimExpansion(GenericOp genericOp,
   // - All the loops are parallel loops.
   return genericOp.hasTensorSemantics() &&
          llvm::all_of(genericOp.getIndexingMaps().getValue(),
                       [](Attribute attr) {
                         return attr.cast<AffineMapAttr>()
                             .getValue()
                             .isProjectedPermutation();
                       }) &&
-         genericOp.getMatchingIndexingMap(fusableOpOperand).getNumResults() > 0 &&
-         llvm::all_of(genericOp.getIteratorTypesArray(), [](StringRef it) {
-           return it == getParallelIteratorTypeName();
-         });
+         genericOp.getMatchingIndexingMap(fusableOpOperand).getNumResults() >
+             0 &&
+         llvm::all_of(genericOp.getIteratorTypesArray(),
+                      utils::isParallelIterator);
 }
 
 namespace {
 /// Information needed to expand a generic operation to fold the reshape with
 /// it.
 class ExpansionInfo {
 public:
   // Computes the mapping from original dimensions of the op to the dimensions
if (expandedOutputType != opOperand->get().getType()) {
           genericOp.getLoc(), expandedOutputType, opOperand->get(),
           reassociation));
     } else {
       outputs.push_back(opOperand->get());
     }
   }
 
   // The iterator types of the expanded op are all parallel.
-  SmallVector<StringRef> iteratorTypes(expansionInfo.getExpandedOpNumDims(),
-                                       getParallelIteratorTypeName());
+  SmallVector<StringRef> iteratorTypes =
+      utils::getNParallelIterators(expansionInfo.getExpandedOpNumDims());
 
   TypeRange resultTypes = ValueRange(outputs).getTypes();
   auto fusedOp =
       rewriter.create<GenericOp>(genericOp.getLoc(), resultTypes,
                                  /*inputs=*/expandedOpOperands, outputs,
                                  expandedOpIndexingMaps, iteratorTypes);
   Region &fusedRegion = fusedOp->getRegion(0);
   Region &originalRegion = genericOp->getRegion(0);
for (ReassociationIndicesRef foldedRangeDims : reassociation) {
     // dims.
     if (llvm::any_of(foldedIterationSpaceDims, [&](int64_t dim) {
           return processedIterationDims.count(dim);
         }))
       continue;
 
     // Check that all folded iterator types are all parallel or all reductions.
     StringRef startIteratorType = iteratorTypes[foldedIterationSpaceDims[0]];
-    if (!isParallelIterator(startIteratorType) &&
-        !isReductionIterator(startIteratorType))
+    if (!utils::isParallelIterator(startIteratorType) &&
+        !utils::isReductionIterator(startIteratorType))
       continue;
     if (llvm::any_of(foldedIterationSpaceDims, [&](int64_t dim) {
           return iteratorTypes[dim] != startIteratorType;
         }))
       continue;
 
     // If the folded dimensions correspond to a "reduction" iterator type,
     // the folded dimensions need to be "in-order". Strictly speaking this is
     // not necessary, for reductions that are associative and commutative,  but
     // using a more strict definition of reduction for now.
-    if (isReductionIterator(startIteratorType)) {
+    if (utils::isReductionIterator(startIteratorType)) {
       bool isContiguous = false;
       for (const auto &startDim : llvm::enumerate(reductionDims)) {
         // Move window in `reductionDims` to start of the folded iteration dims.
         if (startDim.value() != foldedIterationSpaceDims[0])
           continue;
         // If sizes doesnt match, trivial not contiguous. This condition should
         // not be hit.
         if (startDim.index() + foldedIterationSpaceDims.size() >

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

LogicalResult matchAndRewrite(Operation *op,
     if (!isElementwiseMappableOpOnRankedTensors(op))
       return rewriter.notifyMatchFailure(
           op, "requires elementwise op on ranked tensors");
 
     auto rank = op->getResult(0).getType().cast<RankedTensorType>().getRank();
     SmallVector<AffineMap, 3> indexingMaps(
         op->getNumResults() + op->getNumOperands(),
         rewriter.getMultiDimIdentityMap(rank));
-    SmallVector<StringRef, 6> iteratorTypes(rank,
-                                            getParallelIteratorTypeName());
     auto outputs = getOrCreateOperandsMatchingResultTypes(rewriter, op);
     rewriter.replaceOpWithNewOp<linalg::GenericOp>(
         op, /*resultTensorTypes=*/op->getResultTypes(),
         /*inputs=*/op->getOperands(),
         /*outputs=*/outputs,
         /*indexingMaps=*/indexingMaps,
-        /*iteratorTypes=*/iteratorTypes,
+        /*iteratorTypes=*/utils::getNParallelIterators(rank),
         /*bodyBuilder=*/
         [&](OpBuilder &builder, Location loc, ValueRange regionArgs) {
           auto resultTypes = llvm::to_vector<6>(
               llvm::map_range(op->getResultTypes(), [](Type type) {
                 return type.cast<TensorType>().getElementType();
               }));
           auto *scalarOp =
               builder.create(loc, op->getName().getIdentifier(),

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

tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ArrayRef<OpFoldResult> tileSizes,
   SmallVector<linalg::ProcInfo> procInfo;
   if (options.distribution) {
     procInfo.resize(
         iteratorTypes.size(),
         linalg::ProcInfo{nullptr, nullptr, linalg::DistributionMethod::None});
     // Collect loop ranges of tiled loopss, loops that are parallel.
     SmallVector<Range> parallelLoopRanges;
     for (const auto &iteratorType : llvm::enumerate(iteratorTypes)) {
-      if (!isParallelIterator(iteratorType.value()))
+      if (!utils::isParallelIterator(iteratorType.value()))
         break;
       parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
     }
     auto returnedProcInfo =
         options.distribution->procInfo(b, op.getLoc(), parallelLoopRanges);
     unsigned procIdIdx = 0;
     // Update the distribution information for the loops.
     for (const auto &iteratorType : llvm::enumerate(iteratorTypes)) {
-      if (!isParallelIterator(iteratorType.value()))
+      if (!utils::isParallelIterator(iteratorType.value()))
         break;
       procInfo[iteratorType.index()] = returnedProcInfo[procIdIdx++];
     }
   }
 
   // 2. Create the tiled loops.
   LinalgOp res = op;
   SmallVector<Value, 4> ivs, tensorResults;

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

mlir::linalg::LinalgPaddingPattern::returningMatchAndRewrite(
   return paddedOp;
 }
 
 LogicalResult mlir::linalg::CopyVectorizationPattern::matchAndRewrite(
     memref::CopyOp copyOp, PatternRewriter &rewriter) const {
   return vectorizeCopy(rewriter, copyOp);
 }
 
-static SmallVector<StringRef> getNParallelLoopsAttrs(unsigned nParallelLoops) {
-  return SmallVector<StringRef>(nParallelLoops, getParallelIteratorTypeName());
-}
-
 /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp (to
 /// initialize with pad_val) and GenericOp (to copy contents).
 LogicalResult
 PadOpTransformationPattern::matchAndRewrite(tensor::PadOp padOp,
                                             PatternRewriter &rewriter) const {
 
   auto inputShapedType = padOp.getSource().getType().cast<ShapedType>();
   auto resultShapedType = padOp.getResult().getType().cast<ShapedType>();
PadOpTransformationPattern::matchAndRewrite(tensor::PadOp padOp,
 
   SmallVector<AffineMap, 2> transferMaps = {
       rewriter.getMultiDimIdentityMap(inputShapedType.getRank()),
       AffineMap::get(resultShapedType.getRank(),
                      /*symbolCount=*/0, outputExprs, rewriter.getContext())};
 
   rewriter.replaceOpWithNewOp<linalg::GenericOp>(
       padOp, resultShapedType, padOp.getSource(), tmpTensor, transferMaps,
-      getNParallelLoopsAttrs(resultShapedType.getRank()),
+      utils::getNParallelIterators(resultShapedType.getRank()),
       [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
         nestedBuilder.create<linalg::YieldOp>(nestedLoc, args[0]);
       });
 
   return success();
 }
 
 /// Filling `dest` using FillOp constant padding value if possible.

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

static Operation *buildMultiDimReduce(OpBuilder &b, Operation *reduceOp,
                                       const SmallVector<bool> &reductionMask) {
   auto maybeKind = getCombinerOpKind(reduceOp);
   assert(maybeKind && "Failed precondition: could not get reduction kind");
   return b.create<vector::MultiDimReductionOp>(
       reduceOp->getLoc(), valueToReduce, acc, reductionMask, *maybeKind);
 }
 
 static SmallVector<bool> getReductionMask(LinalgOp linalgOp) {
-  return llvm::to_vector(
-      llvm::map_range(linalgOp.getIteratorTypesArray(), isReductionIterator));
+  return llvm::to_vector(llvm::map_range(linalgOp.getIteratorTypesArray(),
+                                         utils::isReductionIterator));
 }
 
 /// Build a vector.transfer_write of `value` into `outputOperand` at indices set
 /// to all `0`; where `outputOperand` is an output operand of the LinalgOp
 /// currently being vectorized. If `dest` has null rank, build an memref.store.
 /// Return the produced value or null if no value is produced.
 static Value buildVectorWrite(OpBuilder &b, Value value,
                               OpOperand *outputOperand) {
vectorizeAsLinalgGeneric(OpBuilder &b, LinalgOp linalgOp,
   }
 
   return success();
 }
 
 // TODO: probably need some extra checks for reduction followed by consumer
 // ops that may not commute (e.g. linear reduction + non-linear instructions).
 static LogicalResult reductionPreconditions(LinalgOp op) {
-  if (llvm::none_of(op.getIteratorTypesArray(), isReductionIterator)) {
+  if (llvm::none_of(op.getIteratorTypesArray(), utils::isReductionIterator)) {
     LDBG("reduction precondition failed: no reduction iterator");
     return failure();
   }
   for (OpOperand *opOperand : op.getDpsInitOperands()) {
     AffineMap indexingMap = op.getMatchingIndexingMap(opOperand);
     if (indexingMap.isPermutation())
       continue;
 

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

bool isElementwise(LinalgOp op) {
   // TODO: relax the restrictions on indexing map.
   for (OpOperand *opOperand : op.getDpsInitOperands()) {
     if (!op.getMatchingIndexingMap(opOperand).isPermutation())
       return false;
   }
   return hasOnlyScalarElementwiseOp(op->getRegion(0));
 }
 
-bool isParallelIterator(StringRef iteratorType) {
-  return iteratorType == getParallelIteratorTypeName();
-}
-
-bool isReductionIterator(StringRef iteratorType) {
-  return iteratorType == getReductionIteratorTypeName();
-}
-
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
     return b.createOrFold<memref::DimOp>(loc, source, dim);
   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
     return b.createOrFold<tensor::DimOp>(loc, source, dim);
   llvm_unreachable("Expected MemRefType or TensorType");
GenericOp makeTransposeOp(OpBuilder &b, Location loc, Value inputTensor,
 
   // Compute the transpose and the indentity indexing maps.
   SmallVector<AffineMap> indexingMaps = {
       inversePermutation(AffineMap::getPermutationMap(
           SmallVector<unsigned>(transposeVector.begin(), transposeVector.end()),
           b.getContext())),
       AffineMap::getMultiDimIdentityMap(transposeVector.size(),
                                         b.getContext())};
-  SmallVector<llvm::StringRef> iteratorTypes(transposeVector.size(),
-                                             getParallelIteratorTypeName());
 
   // Create a GenericOp to transpose `inputTensor` into `outputTensor`.
   auto transposeOp = b.create<GenericOp>(
-      loc, resultTensorType, inputTensor, outputTensor,
-      b.getAffineMapArrayAttr(indexingMaps), b.getStrArrayAttr(iteratorTypes),
-      /*doc=*/nullptr,
-      /*library_call=*/nullptr);
+      loc, resultTensorType, inputTensor, outputTensor, indexingMaps,
+      /*iteratorTypes=*/utils::getNParallelIterators(transposeVector.size()));
   Region &body = transposeOp.getRegion();
   body.push_back(new Block());
   body.front().addArguments({elementType, elementType}, {loc, loc});
 
   // Create the body of the transpose operation.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPointToEnd(&body.front());
   b.create<YieldOp>(loc, transposeOp.getRegion().front().getArgument(0));
   return transposeOp;
 }
 
 GenericOp makeMemRefCopyOp(OpBuilder &b, Location loc, Value from, Value to) {
   auto memrefTypeTo = to.getType().cast<MemRefType>();
 #ifndef NDEBUG
   auto memrefTypeFrom = from.getType().cast<MemRefType>();
   assert(memrefTypeFrom.getRank() == memrefTypeTo.getRank() &&
          "`from` and `to` memref must have the same rank");
 #endif // NDEBUG
 
   AffineMap id =
       AffineMap::getMultiDimIdentityMap(memrefTypeTo.getRank(), b.getContext());
-  SmallVector<StringRef> iteratorTypes(memrefTypeTo.getRank(),
-                                       getParallelIteratorTypeName());
   return b.create<linalg::GenericOp>(
       loc,
       /*inputs=*/from,
       /*outputs=*/to,
       /*indexingMaps=*/llvm::makeArrayRef({id, id}),
-      /*iteratorTypes=*/iteratorTypes,
+      /*iteratorTypes=*/utils::getNParallelIterators(memrefTypeTo.getRank()),
       [](OpBuilder &b, Location loc, ValueRange args) {
         b.create<linalg::YieldOp>(loc, args.front());
       });
 }
 
 /// Specialization to build an scf "for" nest.
 template <>
 void GenerateLoopNest<scf::ForOp>::doit(
static void generateParallelLoopNest(
   // done with it.
   if (iteratorTypes.empty()) {
     bodyBuilderFn(b, loc, ivStorage);
     return;
   }
 
   // If there are no outer parallel loops, generate one sequential loop and
   // recurse.
-  if (!isParallelIterator(iteratorTypes.front())) {
+  if (!utils::isParallelIterator(iteratorTypes.front())) {
     LoopNest singleLoop = buildLoopNest(
         b, loc, lbs.take_front(), ubs.take_front(), steps.take_front(),
         [&](OpBuilder &b, Location loc, ValueRange ivs) {
           ivStorage.append(ivs.begin(), ivs.end());
           generateParallelLoopNest(
               b, loc, lbs.drop_front(), ubs.drop_front(), steps.drop_front(),
               iteratorTypes.drop_front(),
               procInfo.empty() ? procInfo : procInfo.drop_front(),
               bodyBuilderFn, ivStorage);
         });
     return;
   }
 
   unsigned nLoops = iteratorTypes.size();
   unsigned numProcessed = 0;
   DistributionMethod distributionMethod = DistributionMethod::None;
   if (procInfo.empty()) {
-    numProcessed = nLoops - iteratorTypes.drop_while(isParallelIterator).size();
+    numProcessed =
+        nLoops - iteratorTypes.drop_while(utils::isParallelIterator).size();
   } else {
     distributionMethod = procInfo.front().distributionMethod;
     numProcessed =
         nLoops - procInfo
                      .drop_while([&](linalg::ProcInfo p) {
                        return p.distributionMethod == distributionMethod;
                      })
                      .size();

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

 static bool topSortOptimal(unsigned n, ArrayRef<StringRef> iteratorTypes,
                            std::vector<unsigned> &topSort,
                            std::vector<unsigned> &inDegree,
                            std::vector<std::vector<bool>> &adjM) {
   std::vector<unsigned> redIt; // reduce iterator with 0 degree
   std::vector<unsigned> parIt; // parallel iterator with 0 degree
   for (unsigned i = 0; i < n; i++) {
     if (inDegree[i] == 0) {
-      if (linalg::isReductionIterator(iteratorTypes[i]))
+      if (utils::isReductionIterator(iteratorTypes[i]))
         redIt.push_back(i);
       else
         parIt.push_back(i);
     }
   }
 
   while (!redIt.empty() || !parIt.empty()) {
     // We always choose parallel iterator if there is any.
     auto &it = !parIt.empty() ? parIt : redIt;
     auto src = it.back();
     topSort.push_back(src);
     it.pop_back();
     // Update in-degree, and push 0-degree node into worklist.
     for (unsigned dst = 0; dst < n; dst++)
       if (adjM[src][dst] && --inDegree[dst] == 0) {
-        if (linalg::isReductionIterator(iteratorTypes[dst]))
+        if (utils::isReductionIterator(iteratorTypes[dst]))
           redIt.push_back(dst);
         else
           parIt.push_back(dst);
       }
   }
   return topSort.size() == n;
 }
 
static bool isAdmissibleTensorExp(Merger &merger, linalg::GenericOp op,
   // [Bik96,Ch9], is also admissible without special codegen.
   if (merger.isSingleCondition(tensor, exp))
     return true;
   // Accept "truly dynamic" if the output tensor materializes uninitialized
   // into the computation and insertions occur in lexicographic index order.
   if (isMaterializing(lhs->get())) {
     unsigned nest = 0;
     for (unsigned i = 0; i < numLoops; i++) {
-      if (linalg::isReductionIterator(iteratorTypes[topSort[i]]))
+      if (utils::isReductionIterator(iteratorTypes[topSort[i]]))
         break; // terminate at first reduction
       nest++;
     }
     // Determine admissible dynamic insertion situations:
     // (1) fully injective, since there are no reductions,
     // (2) admissible 1-d expansion in innermost dimension.
     if (nest >= op.getRank(lhs) - 1) {
       *sparseOut = lhs;

 /// Generates a for-loop on a single index.
 static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, bool isOuter, bool isInner,
                          unsigned idx, size_t tid, size_t dim,
                          ArrayRef<size_t> extraTids,
                          ArrayRef<size_t> extraDims) {
   Location loc = op.getLoc();
   auto iteratorTypes = op.getIteratorTypesArray();
-  bool isReduction = linalg::isReductionIterator(iteratorTypes[idx]);
+  bool isReduction = utils::isReductionIterator(iteratorTypes[idx]);
   bool isSparse = isCompressedDLT(merger.getDimLevelType(tid, idx)) ||
                   isSingletonDLT(merger.getDimLevelType(tid, idx));
   bool isParallel = isParallelFor(codegen, isOuter, isReduction, isSparse);
   assert(!isParallel);
 
   // Emit a sequential for loop.
   SmallVector<Value, 4> operands;
   if (codegen.redVal)

mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tosa/Utils/CoversionUtils.h"
 
 using namespace mlir;
 using namespace mlir::tosa;
 
-SmallVector<StringRef>
-mlir::tosa::getNParallelLoopsAttrs(unsigned nParallelLoops) {
-  return SmallVector<StringRef>(nParallelLoops, getParallelIteratorTypeName());
-}
-
 SmallVector<Value>
 mlir::tosa::condenseValues(const SmallVector<Value> &values) {
   SmallVector<Value> condensedValues;
   for (auto value : values)
     if (value)
       condensedValues.push_back(value);
   return condensedValues;
 }

mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp

bool mlir::isRowMajorBatchMatmul(ArrayAttr indexingMaps) {
   AffineExpr k = map0.getResult(2);
   auto *context = indexingMaps.getContext();
   auto mapA = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, k}, context));
   auto mapB = AffineMapAttr::get(AffineMap::get(4, 0, {b, k, n}, context));
   auto mapC = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, n}, context));
   auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
   return indexingMaps == maps;
 }
+
+SmallVector<StringRef> mlir::utils::getNParallelIterators(unsigned nParallel) {
+  return SmallVector<StringRef>(nParallel, getParallelIteratorTypeName());
+}