diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h --- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h +++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h @@ -41,12 +41,6 @@ /// 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); diff --git a/mlir/include/mlir/Dialect/Tosa/Utils/CoversionUtils.h b/mlir/include/mlir/Dialect/Tosa/Utils/CoversionUtils.h --- a/mlir/include/mlir/Dialect/Tosa/Utils/CoversionUtils.h +++ b/mlir/include/mlir/Dialect/Tosa/Utils/CoversionUtils.h @@ -21,9 +21,6 @@ namespace mlir { namespace tosa { -// Creates a SmallVector of Stringrefs for N parallel loops -SmallVector getNParallelLoopsAttrs(unsigned nParallelLoops); - // Takes a vector of values and condenses them to a vector with no gaps. SmallVector condenseValues(const SmallVector &values); diff --git a/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h b/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h --- a/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h +++ b/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h @@ -122,6 +122,24 @@ } } +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 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. diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp --- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp +++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp @@ -625,7 +625,7 @@ bool didEncounterError = false; auto linalgOp = rewriter.create( 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()), @@ -1119,7 +1119,7 @@ rewriter.replaceOpWithNewOp( op, resultTy, op.getInput1(), ValueRange{emptyTensor}, affineMaps, - getNParallelLoopsAttrs(resultTy.getRank()), + utils::getNParallelIterators(resultTy.getRank()), [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { nestedBuilder.create(loc, *args.begin()); }); @@ -1230,7 +1230,7 @@ auto linalgOp = rewriter.create( loc, outputTy, genericInputs, ValueRange{emptyTensor}, indexingMaps, - getNParallelLoopsAttrs(rank), + utils::getNParallelIterators(rank), [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) { Value value = blockArgs[0]; @@ -1452,7 +1452,7 @@ Value resize = input; auto genericOp = rewriter.create( loc, resultTy, ValueRange({}), ValueRange{emptyTensor}, affineMaps, - getNParallelLoopsAttrs(resultTy.getRank())); + utils::getNParallelIterators(resultTy.getRank())); resize = genericOp.getResult(0); OpBuilder::InsertionGuard regionGuard(rewriter); @@ -1840,7 +1840,7 @@ rewriter.replaceOpWithNewOp( op, resultTy, ArrayRef({}), ValueRange{emptyTensor}, affineMaps, - getNParallelLoopsAttrs(resultTy.getRank()), + utils::getNParallelIterators(resultTy.getRank()), [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { llvm::SmallVector indices; for (unsigned int i = 0; i < inputTy.getRank(); i++) { @@ -1921,7 +1921,7 @@ auto genericOp = rewriter.create( loc, RankedTensorType::get(genericShape, elementTy), input, ValueRange{emptyTensor}, affineMaps, - getNParallelLoopsAttrs(genericShape.size()), + utils::getNParallelIterators(genericShape.size()), [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { nestedBuilder.create(op.getLoc(), *args.begin()); }); @@ -2083,8 +2083,8 @@ // We need to reduce along the arg-max axis, with parallel operations along // the rest. - SmallVector iteratorTypes; - iteratorTypes.resize(inputTy.getRank(), getParallelIteratorTypeName()); + SmallVector iteratorTypes = + utils::getNParallelIterators(inputTy.getRank()); iteratorTypes[axis] = getReductionIteratorTypeName(); SmallVector srcExprs; @@ -2176,7 +2176,7 @@ auto genericOp = rewriter.create( loc, ArrayRef({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(loc, 0); @@ -2231,7 +2231,7 @@ auto genericOp = rewriter.create( loc, resultTy, ValueRange({input}), ValueRange{emptyTensor}, affineMaps, - getNParallelLoopsAttrs(resultTy.getRank())); + utils::getNParallelIterators(resultTy.getRank())); rewriter.replaceOp(op, genericOp.getResult(0)); { diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp --- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp +++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp @@ -289,7 +289,8 @@ rewriter .create( 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( @@ -311,7 +312,7 @@ rewriter .create( 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( @@ -452,7 +453,7 @@ .create( loc, resultTy, ValueRange({bias, convReshape}), biasEmptyTensor, indexingMaps, - getNParallelLoopsAttrs(resultRank), + utils::getNParallelIterators(resultRank), [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { Value added = nestedBuilder.create( @@ -479,7 +480,7 @@ .create( loc, resultTy, ValueRange({bias, convReshape}), biasEmptyTensor, indexingMaps, - getNParallelLoopsAttrs(resultRank), + utils::getNParallelIterators(resultRank), [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { Value added = nestedBuilder.create( @@ -637,7 +638,8 @@ rewriter .create( 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( @@ -665,7 +667,7 @@ rewriter .create( 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( @@ -825,7 +827,7 @@ loc, ArrayRef({resultTy}), ValueRange{poolingOp}, ValueRange{genericEmptyTensor}, ArrayRef({affineMap, affineMap}), - getNParallelLoopsAttrs(resultTy.getRank()), + utils::getNParallelIterators(resultTy.getRank()), [&](OpBuilder &b, Location loc, ValueRange args) { auto zero = rewriter.create(loc, 0); auto one = rewriter.create(loc, 1); diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp --- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp @@ -321,7 +321,7 @@ 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; @@ -329,7 +329,7 @@ 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; @@ -338,7 +338,7 @@ !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; @@ -346,7 +346,7 @@ 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; @@ -364,7 +364,7 @@ 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; @@ -374,7 +374,7 @@ 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; diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -1419,7 +1419,7 @@ SmallVector MapOp::getIteratorTypesArray() { int64_t rank = getInit().getType().getRank(); - return SmallVector(rank, getParallelIteratorTypeName()); + return utils::getNParallelIterators(rank); } ArrayAttr MapOp::getIndexingMaps() { @@ -1476,8 +1476,8 @@ SmallVector ReduceOp::getIteratorTypesArray() { int64_t inputRank = getInputs()[0].getType().cast().getRank(); - SmallVector iteratorTypes(inputRank, - getParallelIteratorTypeName()); + SmallVector iteratorTypes = + utils::getNParallelIterators(inputRank); for (int64_t reductionDim : getDimensions()) iteratorTypes[reductionDim] = getReductionIteratorTypeName(); return iteratorTypes; @@ -1759,7 +1759,7 @@ SmallVector TransposeOp::getIteratorTypesArray() { int64_t rank = getInit().getType().getRank(); - return SmallVector(rank, getParallelIteratorTypeName()); + return utils::getNParallelIterators(rank); } ArrayAttr TransposeOp::getIndexingMaps() { diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp @@ -470,10 +470,10 @@ .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 { @@ -783,8 +783,8 @@ } // The iterator types of the expanded op are all parallel. - SmallVector iteratorTypes(expansionInfo.getExpandedOpNumDims(), - getParallelIteratorTypeName()); + SmallVector iteratorTypes = + utils::getNParallelIterators(expansionInfo.getExpandedOpNumDims()); TypeRange resultTypes = ValueRange(outputs).getTypes(); auto fusedOp = @@ -1084,8 +1084,8 @@ // 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; @@ -1096,7 +1096,7 @@ // 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. diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp @@ -91,15 +91,13 @@ SmallVector indexingMaps( op->getNumResults() + op->getNumOperands(), rewriter.getMultiDimIdentityMap(rank)); - SmallVector iteratorTypes(rank, - getParallelIteratorTypeName()); auto outputs = getOrCreateOperandsMatchingResultTypes(rewriter, op); rewriter.replaceOpWithNewOp( 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>( diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp @@ -472,7 +472,7 @@ // Collect loop ranges of tiled loopss, loops that are parallel. SmallVector parallelLoopRanges; for (const auto &iteratorType : llvm::enumerate(iteratorTypes)) { - if (!isParallelIterator(iteratorType.value())) + if (!utils::isParallelIterator(iteratorType.value())) break; parallelLoopRanges.push_back(loopRanges[iteratorType.index()]); } @@ -481,7 +481,7 @@ 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++]; } diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @@ -297,10 +297,6 @@ return vectorizeCopy(rewriter, copyOp); } -static SmallVector getNParallelLoopsAttrs(unsigned nParallelLoops) { - return SmallVector(nParallelLoops, getParallelIteratorTypeName()); -} - /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp (to /// initialize with pad_val) and GenericOp (to copy contents). LogicalResult @@ -356,7 +352,7 @@ rewriter.replaceOpWithNewOp( padOp, resultShapedType, padOp.getSource(), tmpTensor, transferMaps, - getNParallelLoopsAttrs(resultShapedType.getRank()), + utils::getNParallelIterators(resultShapedType.getRank()), [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { nestedBuilder.create(nestedLoc, args[0]); }); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -190,8 +190,8 @@ } static SmallVector 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 @@ -611,7 +611,7 @@ // 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(); } diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp --- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp @@ -186,14 +186,6 @@ 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) { @@ -422,15 +414,11 @@ b.getContext())), AffineMap::getMultiDimIdentityMap(transposeVector.size(), b.getContext())}; - SmallVector iteratorTypes(transposeVector.size(), - getParallelIteratorTypeName()); // Create a GenericOp to transpose `inputTensor` into `outputTensor`. auto transposeOp = b.create( - 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}); @@ -452,14 +440,12 @@ AffineMap id = AffineMap::getMultiDimIdentityMap(memrefTypeTo.getRank(), b.getContext()); - SmallVector iteratorTypes(memrefTypeTo.getRank(), - getParallelIteratorTypeName()); return b.create( 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(loc, args.front()); }); @@ -582,7 +568,7 @@ // 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) { @@ -600,7 +586,8 @@ 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 = diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp @@ -186,7 +186,7 @@ std::vector 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); @@ -202,7 +202,7 @@ // 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); @@ -348,7 +348,7 @@ 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++; } @@ -898,7 +898,7 @@ ArrayRef 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); diff --git a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp --- a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp +++ b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp @@ -15,11 +15,6 @@ using namespace mlir; using namespace mlir::tosa; -SmallVector -mlir::tosa::getNParallelLoopsAttrs(unsigned nParallelLoops) { - return SmallVector(nParallelLoops, getParallelIteratorTypeName()); -} - SmallVector mlir::tosa::condenseValues(const SmallVector &values) { SmallVector condensedValues; diff --git a/mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp b/mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp --- a/mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp +++ b/mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp @@ -92,3 +92,7 @@ auto maps = ArrayAttr::get(context, {mapA, mapB, mapC}); return indexingMaps == maps; } + +SmallVector mlir::utils::getNParallelIterators(unsigned nParallel) { + return SmallVector(nParallel, getParallelIteratorTypeName()); +}