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 @@ -42,6 +42,14 @@ inVec = auxVec; } +/// 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); + +/// Given an operation, retrieves the value of each dynamic dimension through +/// constructing the necessary DimOp operators. +SmallVector getDynOperands(Location loc, Value val, OpBuilder &b); + /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp /// is a constant then return a new value set to the smallest such constant. /// If `size` comes from a ConstantOp, return the constant. diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h @@ -30,14 +30,6 @@ /// with `b` at location `loc`. SmallVector getOrCreateRanges(OffsetSizeAndStrideOpInterface op, OpBuilder &b, Location loc); - -/// Given an operation, retrieves the value of each dynamic dimension through -/// constructing the necessary DimOp operators. -SmallVector getDynOperands(Location loc, Value val, OpBuilder &b); - -// 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); } // namespace mlir //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/Vector/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/VectorUtils.h --- a/mlir/include/mlir/Dialect/Vector/VectorUtils.h +++ b/mlir/include/mlir/Dialect/Vector/VectorUtils.h @@ -31,6 +31,10 @@ namespace vector { class TransferWriteOp; class TransferReadOp; + +/// 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); } // namespace vector /// Return the number of elements of basis, `0` if empty. diff --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp --- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp +++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp @@ -166,7 +166,7 @@ Location loc = xferOp.getLoc(); ImplicitLocOpBuilder lb(xferOp.getLoc(), b); if (!xferOp.isDimInBounds(0) && !isBroadcast) { - Value memrefDim = createOrFoldDimOp(b, loc, xferOp.source(), *dim); + Value memrefDim = vector::createOrFoldDimOp(b, loc, xferOp.source(), *dim); AffineExpr d0, d1; bindDims(xferOp.getContext(), d0, d1); Value base = xferOp.indices()[dim.getValue()]; 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 @@ -191,6 +191,17 @@ return res; } +/// Helper function that creates a memref::DimOp or tensor::DimOp depending on +/// the type of `source`. +static Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, + int64_t dim) { + if (source.getType().isa()) + return b.createOrFold(loc, source, dim); + if (source.getType().isa()) + return b.createOrFold(loc, source, dim); + llvm_unreachable("Expected MemRefType or TensorType"); +} + SmallVector LinalgOp::createFlatListOfOperandDims(OpBuilder &b, Location loc) { SmallVector res; diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp @@ -108,6 +108,7 @@ #include "PassDetail.h" #include "mlir/Dialect/Linalg/IR/LinalgOps.h" #include "mlir/Dialect/Linalg/Passes.h" +#include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/SCF.h" #include "mlir/Dialect/Utils/StaticValueUtils.h" 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 @@ -66,7 +66,7 @@ Value firstOperand = operands.front(); auto rankedTensorType = t.cast(); auto staticShape = llvm::to_vector<4>(rankedTensorType.getShape()); - auto dynamicShape = getDynOperands(loc, firstOperand, b); + auto dynamicShape = linalg::getDynOperands(loc, firstOperand, b); res.push_back(b.create( loc, dynamicShape, staticShape, rankedTensorType.getElementType())); 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 @@ -157,6 +157,28 @@ namespace mlir { namespace linalg { +/// 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()) + return b.createOrFold(loc, source, dim); + if (source.getType().isa()) + return b.createOrFold(loc, source, dim); + llvm_unreachable("Expected MemRefType or TensorType"); +} + +/// Given an operation, retrieves the value of each dynamic dimension through +/// constructing the necessary DimOp operators. +SmallVector getDynOperands(Location loc, Value val, OpBuilder &b) { + SmallVector dynOperands; + auto shapedType = val.getType().cast(); + for (auto dim : llvm::enumerate(shapedType.getShape())) { + if (dim.value() == ShapedType::kDynamicSize) + dynOperands.push_back(createOrFoldDimOp(b, loc, val, dim.index())); + } + return dynOperands; +} + /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp /// is a constant then return a new value set to the smallest such constant. /// Otherwise returngetSmallestBoundingIndex nullptr. diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp --- a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp +++ b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp @@ -33,29 +33,6 @@ }; } // end anonymous namespace -SmallVector mlir::getDynOperands(Location loc, Value val, - OpBuilder &b) { - SmallVector dynOperands; - auto shapedType = val.getType().cast(); - for (auto dim : llvm::enumerate(shapedType.getShape())) { - if (dim.value() == ShapedType::kDynamicSize) - dynOperands.push_back(createOrFoldDimOp(b, loc, val, dim.index())); - } - return dynOperands; -} - -// Helper function that creates a memref::DimOp or tensor::DimOp depending on -// the type of `source`. -// TODO: Move helper function out of MemRef dialect. -Value mlir::createOrFoldDimOp(OpBuilder &b, Location loc, Value source, - int64_t dim) { - if (source.getType().isa()) - return b.createOrFold(loc, source, dim); - if (source.getType().isa()) - return b.createOrFold(loc, source, dim); - llvm_unreachable("Expected MemRefType or TensorType"); -} - void mlir::memref::MemRefDialect::initialize() { addOperationsget(), d); + up = rewriter.create(loc, t->get(), d); args.push_back(up); } else { up = rewriter.create(loc, shape[d]); diff --git a/mlir/lib/Dialect/Vector/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/VectorTransforms.cpp --- a/mlir/lib/Dialect/Vector/VectorTransforms.cpp +++ b/mlir/lib/Dialect/Vector/VectorTransforms.cpp @@ -2330,7 +2330,7 @@ Value sum = makeComposedAffineApply(b, loc, d0 + vs, xferOp.indices()[indicesIdx]); Value cond = createFoldedSLE( - b, sum, createOrFoldDimOp(b, loc, xferOp.source(), indicesIdx)); + b, sum, vector::createOrFoldDimOp(b, loc, xferOp.source(), indicesIdx)); if (!cond) return; // Conjunction over all dims for which we are in-bounds. @@ -2415,8 +2415,8 @@ auto isaWrite = isa(xferOp); xferOp.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) { using MapList = ArrayRef>; - Value dimMemRef = - createOrFoldDimOp(b, xferOp.getLoc(), xferOp.source(), indicesIdx); + Value dimMemRef = vector::createOrFoldDimOp(b, xferOp.getLoc(), + xferOp.source(), indicesIdx); Value dimAlloc = lb.create(alloc, resultIdx); Value index = xferOp.indices()[indicesIdx]; AffineExpr i, j, k; @@ -3954,7 +3954,8 @@ unsigned vecWidth = vtp.getNumElements(); unsigned lastIndex = llvm::size(xferOp.indices()) - 1; Value off = xferOp.indices()[lastIndex]; - Value dim = createOrFoldDimOp(rewriter, loc, xferOp.source(), lastIndex); + Value dim = + vector::createOrFoldDimOp(rewriter, loc, xferOp.source(), lastIndex); Value mask = buildVectorComparison( rewriter, xferOp, enableIndexOptimizations, vecWidth, dim, &off); diff --git a/mlir/lib/Dialect/Vector/VectorUtils.cpp b/mlir/lib/Dialect/Vector/VectorUtils.cpp --- a/mlir/lib/Dialect/Vector/VectorUtils.cpp +++ b/mlir/lib/Dialect/Vector/VectorUtils.cpp @@ -13,7 +13,9 @@ #include "mlir/Dialect/Vector/VectorUtils.h" #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/IR/Builders.h" #include "mlir/IR/IntegerSet.h" @@ -27,6 +29,17 @@ using namespace mlir; +/// Helper function that creates a memref::DimOp or tensor::DimOp depending on +/// the type of `source`. +Value mlir::vector::createOrFoldDimOp(OpBuilder &b, Location loc, Value source, + int64_t dim) { + if (source.getType().isa()) + return b.createOrFold(loc, source, dim); + if (source.getType().isa()) + return b.createOrFold(loc, source, dim); + llvm_unreachable("Expected MemRefType or TensorType"); +} + /// Return the number of elements of basis, `0` if empty. int64_t mlir::computeMaxLinearIndex(ArrayRef basis) { if (basis.empty()) diff --git a/mlir/lib/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Transforms/PipelineDataTransfer.cpp --- a/mlir/lib/Transforms/PipelineDataTransfer.cpp +++ b/mlir/lib/Transforms/PipelineDataTransfer.cpp @@ -85,8 +85,12 @@ // The double buffer is allocated right before 'forOp'. OpBuilder bOuter(forOp); // Put together alloc operands for any dynamic dimensions of the memref. - - auto allocOperands = getDynOperands(forOp.getLoc(), oldMemRef, bOuter); + SmallVector allocOperands; + for (auto dim : llvm::enumerate(oldMemRefType.getShape())) { + if (dim.value() == ShapedType::kDynamicSize) + allocOperands.push_back(bOuter.createOrFold( + forOp.getLoc(), oldMemRef, dim.index())); + } // Create and place the alloc right before the 'affine.for' operation. Value newMemRef = bOuter.create(