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 @@ -187,15 +187,6 @@ LinalgOp op, ValueRange operands, ValueRange results); -/// Turns an OpFoldResult into a value, creating an index-typed constant if -/// necessary. -Value materializeOpFoldResult(ImplicitLocOpBuilder &builder, - OpFoldResult opFoldResult); -Value materializeOpFoldResult(OpBuilder &b, Location loc, - OpFoldResult opFoldResult); -Value materializeOpFoldResult(OpBuilder &b, Location loc, - ArrayRef opFoldResults); - /// A struct containg offsets-sizes-strides arguments of the tiled shape. struct SliceParameters { SmallVector offsets; 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 @@ -13,9 +13,9 @@ #include "mlir/Dialect/Linalg/Passes.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arithmetic/Utils/Utils.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" @@ -1453,7 +1453,7 @@ rewriter.setInsertionPoint(collapsedGenericOp); SmallVector loopBound = llvm::to_vector(llvm::map_range(loopRanges, [&](Range range) { - return materializeOpFoldResult(rewriter, loc, range.size); + return getValueOrCreateConstantIndexOp(rewriter, loc, range.size); })); generateCollapsedIndexingRegion(loc, &collapsedGenericOp->getRegion(0).front(), diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp @@ -11,11 +11,11 @@ //===----------------------------------------------------------------------===// #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" +#include "mlir/Dialect/Arithmetic/Utils/Utils.h" #include "mlir/Dialect/Complex/IR/Complex.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Passes.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineExprVisitor.h" @@ -227,9 +227,10 @@ LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n"); Value size; if (auto attr = rangeValue.size.dyn_cast()) { - size = materializeOpFoldResult(b, loc, rangeValue.size); + size = getValueOrCreateConstantIndexOp(b, loc, rangeValue.size); } else { - Value materializedSize = materializeOpFoldResult(b, loc, rangeValue.size); + Value materializedSize = + getValueOrCreateConstantIndexOp(b, loc, rangeValue.size); FailureOr upperBound = getConstantUpperBoundForIndex(materializedSize); size = failed(upperBound) 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 @@ -18,7 +18,6 @@ #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/Transforms/Transforms.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" @@ -123,13 +122,15 @@ return failure(); // The code below works only on values. - ImplicitLocOpBuilder b(op.getLoc(), builder); + Location loc = op.getLoc(); + ImplicitLocOpBuilder b(loc, builder); if (emitAssertions) { emitIsPositiveIndexAssertion(b, targetSize); emitIsPositiveIndexAssertion(b, divisor); } - Value targetSizeValue = materializeOpFoldResult(b, targetSize); - Value divisorValue = materializeOpFoldResult(b, divisor); + Value targetSizeValue = + getValueOrCreateConstantIndexOp(builder, loc, targetSize); + Value divisorValue = getValueOrCreateConstantIndexOp(builder, loc, divisor); // Find the trip count of the iteration space dimension for which the tile // sizes are computed. @@ -140,7 +141,7 @@ makeComposedFoldedMultiResultAffineApply(b, op.getLoc(), shapesToLoops, allShapes); Value tripCount = - materializeOpFoldResult(b, op.getLoc(), loopRanges[dimension]); + getValueOrCreateConstantIndexOp(b, op.getLoc(), loopRanges[dimension]); // Compute the tile sizes and the respective numbers of tiles. AffineExpr s0 = b.getAffineSymbolExpr(0); @@ -245,8 +246,7 @@ })); SmallVector materializedNonZeroNumThreads = llvm::to_vector(llvm::map_range(nonZeroNumThreads, [&](OpFoldResult ofr) { - ImplicitLocOpBuilder ilocb(loc, b); - return materializeOpFoldResult(ilocb, ofr); + return getValueOrCreateConstantIndexOp(b, loc, ofr); })); Operation *tiledOp = nullptr; @@ -618,9 +618,12 @@ for (int64_t i = 0; i < rank; ++i) { allDims.push_back(ranges[i].size); if (!isZero(tileSizes[i])) { - lbs.push_back(materializeOpFoldResult(builder, loc, ranges[i].offset)); - dims.push_back(materializeOpFoldResult(builder, loc, ranges[i].size)); - steps.push_back(materializeOpFoldResult(builder, loc, tileSizes[i])); + lbs.push_back( + getValueOrCreateConstantIndexOp(builder, loc, ranges[i].offset)); + dims.push_back( + getValueOrCreateConstantIndexOp(builder, loc, ranges[i].size)); + steps.push_back( + getValueOrCreateConstantIndexOp(builder, loc, tileSizes[i])); } } // Generate loop nest: One loop per dimension. 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 @@ -140,9 +140,11 @@ SmallVectorImpl &ubs, SmallVectorImpl &steps) { for (Range range : ranges) { - lbs.emplace_back(materializeOpFoldResult(builder, loc, range.offset)); - ubs.emplace_back(materializeOpFoldResult(builder, loc, range.size)); - steps.emplace_back(materializeOpFoldResult(builder, loc, range.stride)); + lbs.emplace_back( + getValueOrCreateConstantIndexOp(builder, loc, range.offset)); + ubs.emplace_back(getValueOrCreateConstantIndexOp(builder, loc, range.size)); + steps.emplace_back( + getValueOrCreateConstantIndexOp(builder, loc, range.stride)); } } @@ -929,33 +931,6 @@ return tensorResults; } -Value materializeOpFoldResult(ImplicitLocOpBuilder &builder, - OpFoldResult opFoldResult) { - if (!opFoldResult) - return nullptr; - - if (auto value = opFoldResult.dyn_cast()) - return value; - auto attr = opFoldResult.get().cast(); - return builder.create(attr.getValue().getSExtValue()); -} - -Value materializeOpFoldResult(OpBuilder &builder, Location loc, - OpFoldResult opFoldResult) { - ImplicitLocOpBuilder b(loc, builder); - return materializeOpFoldResult(b, opFoldResult); -} - -SmallVector -materializeOpFoldResults(OpBuilder &builder, Location loc, - ArrayRef opFoldResults) { - ImplicitLocOpBuilder b(loc, builder); - SmallVector values; - for (const auto &opFoldResult : opFoldResults) - values.push_back(materializeOpFoldResult(b, opFoldResult)); - return values; -} - SmallVector> computeAllSliceParameters(OpBuilder &builder, Location loc, LinalgOp linalgOp, ValueRange valuesToTile, ArrayRef ivs, @@ -1046,7 +1021,8 @@ OpFoldResult applied = makeComposedFoldedAffineApply( b, indexOp.getLoc(), index + offset, {getAsOpFoldResult(indexOp.getResult()), offsets[indexOp.getDim()]}); - Value materialized = materializeOpFoldResult(b, indexOp.getLoc(), applied); + Value materialized = + getValueOrCreateConstantIndexOp(b, indexOp.getLoc(), applied); b.replaceOpWithIf(indexOp, materialized, [&](OpOperand &use) { return use.getOwner() != materialized.getDefiningOp(); });