[mlir][NFC] MemRef cleanup: Remove helper functions

[springerm]

Summary

Remove getDynOperands and createOrFoldDimOp from MemRef.h to decouple MemRef a bit from Tensor. These two functions are used in other dialects/transforms.

Depends On D105256

[springerm]

rebase

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

 void applyPermutationToVector(SmallVector<T, N> &inVec,
                               ArrayRef<unsigned> permutation) {
   SmallVector<T, N> auxVec(inVec.size());
   for (unsigned i = 0; i < permutation.size(); ++i)
     auxVec[i] = inVec[permutation[i]];
   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<Value, 4> 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.
 /// Otherwise return nullptr.
 IntegerAttr getSmallestBoundingIndex(Value size);
 
 //===----------------------------------------------------------------------===//
 // Iterator type utilities

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

 
 raw_ostream &operator<<(raw_ostream &os, Range &range);
 
 /// Return the list of Range (i.e. offset, size, stride). Each Range
 /// entry contains either the dynamic value or a ConstantIndexOp constructed
 /// with `b` at location `loc`.
 SmallVector<Range, 8> getOrCreateRanges(OffsetSizeAndStrideOpInterface op,
                                         OpBuilder &b, Location loc);
-
-/// Given an operation, retrieves the value of each dynamic dimension through
-/// constructing the necessary DimOp operators.
-SmallVector<Value, 4> 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
 
 //===----------------------------------------------------------------------===//
 // MemRef Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/IR/MemRefOpsDialect.h.inc"
 

mlir/include/mlir/Dialect/Vector/VectorUtils.h

 class ShapedType;
 class Value;
 class VectorType;
 class VectorTransferOpInterface;
 
 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.
 int64_t computeMaxLinearIndex(ArrayRef<int64_t> basis);
 
 /// Given a shape with sizes greater than 0 along all dimensions,
 /// return the distance, in number of elements, between a slice in a dimension
 /// and the next slice in the same dimension.

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

static Value generateInBoundsCheck(
   bool hasRetVal = !resultTypes.empty();
   Value cond; // Condition to be built...
 
   // Condition check 1: Access in-bounds?
   bool isBroadcast = !dim.hasValue(); // No in-bounds check for broadcasts.
   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()];
     Value memrefIdx = makeComposedAffineApply(b, loc, d0 + d1, {base, iv});
     cond = lb.create<CmpIOp>(CmpIPredicate::sgt, memrefDim, memrefIdx);
   }
 
   // Condition check 2: Masked in?

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

SmallVector<Value, 4> mlir::linalg::applyMapToValues(OpBuilder &b, Location loc,
   // folding occurs eagerly.
   for (auto expr : map.getResults()) {
     AffineMap map = AffineMap::get(numDims, numSym, expr);
     res.push_back(createFoldedComposedAffineApply(b, loc, map, values));
   }
   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<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");
+}
+
 SmallVector<Value, 4> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,
                                                             Location loc) {
   SmallVector<Value, 4> res;
   for (OpOperand *opOperand : getInputAndOutputOperands()) {
     for (int64_t i = 0, e = getRank(opOperand); i < e; ++i)
       res.push_back(createOrFoldDimOp(b, loc, opOperand->get(), i));
   }
   return res;

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

 //  Lastly, note that layout map chosen to bufferize is the most dynamic
 //  canonical strided layout of the proper rank. This ensures compatibility with
 //  expected layouts after transformations. Combinations of memref.cast +
 //  canonicalization are responsible for clean ups.
 
 #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"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/BufferUtils.h"

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

for (Type t : rankedTensorTypes) {
     }
     if (found)
       continue;
 
     // Extract static / dynamic shape mix from the first operand.
     Value firstOperand = operands.front();
     auto rankedTensorType = t.cast<RankedTensorType>();
     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<linalg::InitTensorOp>(
         loc, dynamicShape, staticShape, rankedTensorType.getElementType()));
   }
   return res;
 }
 
 namespace {

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

for (Range range : ranges) {
     ubs.emplace_back(range.size);
     steps.emplace_back(range.stride);
   }
 }
 
 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<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");
+}
+
+/// Given an operation, retrieves the value of each dynamic dimension through
+/// constructing the necessary DimOp operators.
+SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b) {
+  SmallVector<Value, 4> dynOperands;
+  auto shapedType = val.getType().cast<ShapedType>();
+  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.
 IntegerAttr getSmallestBoundingIndex(Value size) {
   Optional<int64_t> boundingConst = {};
   if (auto affineMinOp = size.getDefiningOp<AffineMinOp>()) {
     for (auto e : affineMinOp.getAffineMap().getResults())
       if (auto cst = e.dyn_cast<AffineConstantExpr>())

mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp

struct MemRefInlinerInterface : public DialectInlinerInterface {
   }
   bool isLegalToInline(Operation *, Region *, bool wouldBeCloned,
                        BlockAndValueMapping &) const final {
     return true;
   }
 };
 } // end anonymous namespace
 
-SmallVector<Value, 4> mlir::getDynOperands(Location loc, Value val,
-                                           OpBuilder &b) {
-  SmallVector<Value, 4> dynOperands;
-  auto shapedType = val.getType().cast<ShapedType>();
-  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<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");
-}
-
 void mlir::memref::MemRefDialect::initialize() {
   addOperations<DmaStartOp, DmaWaitOp,
 #define GET_OP_LIST
 #include "mlir/Dialect/MemRef/IR/MemRefOps.cpp.inc"
                 >();
   addInterfaces<MemRefInlinerInterface>();
 }

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

for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
         codegen.pointers[tensor][idx] =
             rewriter.create<ToPointersOp>(loc, ptrTp, t->get(), dim);
         codegen.indices[tensor][idx] =
             rewriter.create<ToIndicesOp>(loc, indTp, t->get(), dim);
       }
       // Find lower and upper bound in current dimension.
       Value up;
       if (shape[d] == MemRefType::kDynamicSize) {
-        up = createOrFoldDimOp(rewriter, loc, t->get(), d);
+        up = rewriter.create<tensor::DimOp>(loc, t->get(), d);
         args.push_back(up);
       } else {
         up = rewriter.create<ConstantIndexOp>(loc, shape[d]);
       }
       codegen.sizes[idx] = codegen.highs[tensor][idx] = up;
     }
     // Perform the required bufferization. Dense inputs materialize
     // from the input tensors. Dense outputs need special handling.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

xferOp.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) {
     Location loc = xferOp.getLoc();
     ImplicitLocOpBuilder lb(loc, b);
     int64_t vectorSize = xferOp.getVectorType().getDimSize(resultIdx);
     auto d0 = getAffineDimExpr(0, xferOp.getContext());
     auto vs = getAffineConstantExpr(vectorSize, xferOp.getContext());
     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.
     if (inBoundsCond)
       inBoundsCond = lb.create<AndOp>(inBoundsCond, cond);
     else
       inBoundsCond = cond;
   });
assert(memrefRank == alloc.getType().cast<MemRefType>().getRank() &&
          "Expected memref rank to match the alloc rank");
   ValueRange leadingIndices =
       xferOp.indices().take_front(xferOp.getLeadingShapedRank());
   SmallVector<OpFoldResult, 4> sizes;
   sizes.append(leadingIndices.begin(), leadingIndices.end());
   auto isaWrite = isa<vector::TransferWriteOp>(xferOp);
   xferOp.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) {
     using MapList = ArrayRef<ArrayRef<AffineExpr>>;
-    Value dimMemRef =
-        createOrFoldDimOp(b, xferOp.getLoc(), xferOp.source(), indicesIdx);
+    Value dimMemRef = vector::createOrFoldDimOp(b, xferOp.getLoc(),
+                                                xferOp.source(), indicesIdx);
     Value dimAlloc = lb.create<memref::DimOp>(alloc, resultIdx);
     Value index = xferOp.indices()[indicesIdx];
     AffineExpr i, j, k;
     bindDims(xferOp.getContext(), i, j, k);
     SmallVector<AffineMap, 4> maps =
         AffineMap::inferFromExprList(MapList{{i - j, k}});
     // affine_min(%dimMemRef - %index, %dimAlloc)
     Value affineMin = lb.create<AffineMinOp>(
LogicalResult matchAndRewrite(ConcreteOp xferOp,
     //   (in-bounds mask):
     //   [ offset + 0 .. offset + vector_length - 1 ] < [ dim .. dim ]
     //
     // TODO: when the leaf transfer rank is k > 1, we need the last `k`
     //       dimensions here.
     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);
 
     if (xferOp.mask()) {
       // Intersect the in-bounds with the mask specified as an op parameter.
       mask = rewriter.create<AndOp>(loc, mask, xferOp.mask());
     }
 

mlir/lib/Dialect/Vector/VectorUtils.cpp

 //===- VectorUtils.cpp - MLIR Utilities for VectorOps   ------------------===//
 //
 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utility methods for working with the Vector dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
 #include <numeric>
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 
 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<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");
+}
+
 /// Return the number of elements of basis, `0` if empty.
 int64_t mlir::computeMaxLinearIndex(ArrayRef<int64_t> basis) {
   if (basis.empty())
     return 0;
   return std::accumulate(basis.begin(), basis.end(), 1,
                          std::multiplies<int64_t>());
 }
 

mlir/lib/Transforms/PipelineDataTransfer.cpp

static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
   };
 
   auto oldMemRefType = oldMemRef.getType().cast<MemRefType>();
   auto newMemRefType = doubleShape(oldMemRefType);
 
   // 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<Value, 4> allocOperands;
+  for (auto dim : llvm::enumerate(oldMemRefType.getShape())) {
+    if (dim.value() == ShapedType::kDynamicSize)
+      allocOperands.push_back(bOuter.createOrFold<memref::DimOp>(
+          forOp.getLoc(), oldMemRef, dim.index()));
+  }
 
   // Create and place the alloc right before the 'affine.for' operation.
   Value newMemRef = bOuter.create<memref::AllocOp>(
       forOp.getLoc(), newMemRefType, allocOperands);
 
   // Create 'iv mod 2' value to index the leading dimension.
   auto d0 = bInner.getAffineDimExpr(0);
   int64_t step = forOp.getStep();