MemRefUtils.cpp

Go to the documentation of this file.

//===- MemRefUtils.cpp - Utilities to support the MemRef dialect ----------===//
//
// Part of the LLVM 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 utilities for the MemRef dialect.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
#include "llvm/ADT/STLExtras.h"
 
namespace mlir {
namespace memref {
 
bool isStaticShapeAndContiguousRowMajor(MemRefType type) {
  if (!type.hasStaticShape())
    return false;
 
  int64_t rank = type.getRank();
  if (rank == 0)
    return true;
 
  SmallVector<int64_t> strides;
  int64_t offset;
  if (failed(type.getStridesAndOffset(strides, offset)))
    return false;
 
  // MemRef is contiguous if outer dimensions are size-1 and inner
  // dimensions have unit strides.
  int64_t runningStride = 1;
  int64_t curDim = rank - 1;
  // Finds all inner dimensions with unit strides.
  while (curDim >= 0 && strides[curDim] == runningStride) {
    runningStride *= type.getDimSize(curDim);
    --curDim;
  }
 
  // Check if other dimensions are size-1.
  while (curDim >= 0 && type.getDimSize(curDim) == 1) {
    --curDim;
  }
 
  // All dims are unit-strided or size-1.
  return curDim < 0;
}
 
std::pair<LinearizedMemRefInfo, OpFoldResult> getLinearizedMemRefOffsetAndSize(
    OpBuilder &builder, Location loc, int srcBits, int dstBits,
    OpFoldResult offset, ArrayRef<OpFoldResult> sizes,
    ArrayRef<OpFoldResult> strides, ArrayRef<OpFoldResult> indices,
    LinearizedDivKind sizeDivKind) {
  unsigned sourceRank = sizes.size();
  assert(sizes.size() == strides.size() &&
         "expected as many sizes as strides for a memref");
  SmallVector<OpFoldResult> indicesVec = llvm::to_vector(indices);
  if (indices.empty())
    indicesVec.resize(sourceRank, builder.getIndexAttr(0));
  assert(indicesVec.size() == strides.size() &&
         "expected as many indices as rank of memref");
 
  // Create the affine symbols and values for linearization.
  SmallVector<AffineExpr> symbols(2 * sourceRank);
  bindSymbolsList(builder.getContext(), MutableArrayRef{symbols});
  AffineExpr addMulMap = builder.getAffineConstantExpr(0);
 
  SmallVector<OpFoldResult> offsetValues(2 * sourceRank);
 
  for (unsigned i = 0; i < sourceRank; ++i) {
    unsigned offsetIdx = 2 * i;
    addMulMap = addMulMap + symbols[offsetIdx] * symbols[offsetIdx + 1];
    offsetValues[offsetIdx] = indicesVec[i];
    offsetValues[offsetIdx + 1] = strides[i];
  }
  // Adjust linearizedIndices and size by the scale factor (dstBits / srcBits).
  int64_t scaler = dstBits / srcBits;
  OpFoldResult linearizedIndices = affine::makeComposedFoldedAffineApply(
      builder, loc, addMulMap.floorDiv(scaler), offsetValues);
 
  size_t symbolIndex = 0;
  SmallVector<OpFoldResult> values;
  SmallVector<AffineExpr> productExpressions;
  for (unsigned i = 0; i < sourceRank; ++i) {
    AffineExpr strideExpr = symbols[symbolIndex++];
    values.push_back(strides[i]);
    AffineExpr sizeExpr = symbols[symbolIndex++];
    values.push_back(sizes[i]);
 
    AffineExpr product = strideExpr * sizeExpr;
    productExpressions.push_back(sizeDivKind == LinearizedDivKind::Ceil
                                     ? product.ceilDiv(scaler)
                                     : product.floorDiv(scaler));
  }
  AffineMap maxMap = AffineMap::get(
      /*dimCount=*/0, /*symbolCount=*/symbolIndex, productExpressions,
      builder.getContext());
  OpFoldResult linearizedSize =
      affine::makeComposedFoldedAffineMax(builder, loc, maxMap, values);
 
  // Adjust baseOffset by the scale factor (dstBits / srcBits).
  AffineExpr s0;
  bindSymbols(builder.getContext(), s0);
  OpFoldResult adjustBaseOffset = affine::makeComposedFoldedAffineApply(
      builder, loc, s0.floorDiv(scaler), {offset});
 
  OpFoldResult intraVectorOffset = affine::makeComposedFoldedAffineApply(
      builder, loc, addMulMap % scaler, offsetValues);
 
  return {{adjustBaseOffset, linearizedSize, intraVectorOffset},
          linearizedIndices};
}
 
LinearizedMemRefInfo
getLinearizedMemRefOffsetAndSize(OpBuilder &builder, Location loc, int srcBits,
                                 int dstBits, OpFoldResult offset,
                                 ArrayRef<OpFoldResult> sizes,
                                 LinearizedDivKind sizeDivKind) {
  SmallVector<OpFoldResult> strides(sizes.size());
  if (!sizes.empty()) {
    strides.back() = builder.getIndexAttr(1);
    AffineExpr s0, s1;
    bindSymbols(builder.getContext(), s0, s1);
    for (int64_t index = static_cast<int64_t>(sizes.size()) - 1; index > 0;
         --index) {
      strides[index - 1] = affine::makeComposedFoldedAffineApply(
          builder, loc, s0 * s1,
          ArrayRef<OpFoldResult>{strides[index], sizes[index]});
    }
  }
 
  LinearizedMemRefInfo linearizedMemRefInfo;
  std::tie(linearizedMemRefInfo, std::ignore) =
      getLinearizedMemRefOffsetAndSize(builder, loc, srcBits, dstBits, offset,
                                       sizes, strides, /*indices=*/{},
                                       sizeDivKind);
  return linearizedMemRefInfo;
}
 
/// Returns true if all the uses of op are not read/load.
/// There can be view-like-op users as long as all its users are also
/// StoreOp/transfer_write. If return true it also fills out the uses, if it
/// returns false uses is unchanged.
static bool resultIsNotRead(Operation *op, std::vector<Operation *> &uses) {
  std::vector<Operation *> opUses;
  for (OpOperand &use : op->getUses()) {
    Operation *useOp = use.getOwner();
    // Use escaped the scope
    if (useOp->mightHaveTrait<OpTrait::IsTerminator>())
      return false;
    if (isa<memref::DeallocOp>(useOp) ||
        (useOp->getNumResults() == 0 && useOp->getNumRegions() == 0 &&
         !mlir::hasEffect<MemoryEffects::Read>(useOp)) ||
        (isa<ViewLikeOpInterface>(useOp) && resultIsNotRead(useOp, opUses))) {
      opUses.push_back(useOp);
      continue;
    }
    return false;
  }
  llvm::append_range(uses, opUses);
  return true;
}
 
void eraseDeadAllocAndStores(RewriterBase &rewriter, Operation *parentOp) {
  std::vector<Operation *> opToErase;
  parentOp->walk([&](Operation *op) {
    std::vector<Operation *> candidates;
    if (isa<memref::AllocOp, memref::AllocaOp>(op) &&
        resultIsNotRead(op, candidates)) {
      llvm::append_range(opToErase, candidates);
      opToErase.push_back(op);
    }
  });
 
  for (Operation *op : opToErase)
    rewriter.eraseOp(op);
}
 
static SmallVector<OpFoldResult>
computeSuffixProductIRBlockImpl(Location loc, OpBuilder &builder,
                                ArrayRef<OpFoldResult> sizes,
                                OpFoldResult unit) {
  SmallVector<OpFoldResult> strides(sizes.size(), unit);
  AffineExpr s0, s1;
  bindSymbols(builder.getContext(), s0, s1);
 
  for (int64_t r = static_cast<int64_t>(strides.size()) - 1; r > 0; --r) {
    strides[r - 1] = affine::makeComposedFoldedAffineApply(
        builder, loc, s0 * s1, {strides[r], sizes[r]});
  }
  return strides;
}
 
SmallVector<OpFoldResult>
computeSuffixProductIRBlock(Location loc, OpBuilder &builder,
                            ArrayRef<OpFoldResult> sizes) {
  OpFoldResult unit = builder.getIndexAttr(1);
  return computeSuffixProductIRBlockImpl(loc, builder, sizes, unit);
}
 
MemrefValue skipFullyAliasingOperations(MemrefValue source) {
  while (auto *op = source.getDefiningOp()) {
    if (auto subViewOp = dyn_cast<memref::SubViewOp>(op);
        subViewOp && subViewOp.hasZeroOffset() && subViewOp.hasUnitStride()) {
      // A `memref.subview` with an all zero offset, and all unit strides, still
      // points to the same memory.
      source = cast<MemrefValue>(subViewOp.getSource());
    } else if (auto castOp = dyn_cast<memref::CastOp>(op)) {
      // A `memref.cast` still points to the same memory.
      source = castOp.getSource();
    } else {
      return source;
    }
  }
  return source;
}
 
MemrefValue skipViewLikeOps(MemrefValue source) {
  while (auto *op = source.getDefiningOp()) {
    if (auto viewLike = dyn_cast<ViewLikeOpInterface>(op)) {
      if (source == viewLike.getViewDest()) {
        source = cast<MemrefValue>(viewLike.getViewSource());
        continue;
      }
    }
    return source;
  }
  return source;
}
 
void resolveSourceIndicesExpandShape(Location loc, PatternRewriter &rewriter,
                                     memref::ExpandShapeOp expandShapeOp,
                                     ValueRange indices,
                                     SmallVectorImpl<Value> &sourceIndices,
                                     bool startsInbounds) {
  SmallVector<OpFoldResult> destShape = expandShapeOp.getMixedOutputShape();
 
  // Traverse all reassociation groups to determine the appropriate indices
  // corresponding to each one of them post op folding.
  for (ArrayRef<int64_t> group : expandShapeOp.getReassociationIndices()) {
    assert(!group.empty() && "association indices groups cannot be empty");
    int64_t groupSize = group.size();
    if (groupSize == 1) {
      sourceIndices.push_back(indices[group[0]]);
      continue;
    }
    SmallVector<OpFoldResult> groupBasis =
        llvm::map_to_vector(group, [&](int64_t d) { return destShape[d]; });
    SmallVector<Value> groupIndices =
        llvm::map_to_vector(group, [&](int64_t d) { return indices[d]; });
    Value collapsedIndex = affine::AffineLinearizeIndexOp::create(
        rewriter, loc, groupIndices, groupBasis, /*disjoint=*/startsInbounds);
    sourceIndices.push_back(collapsedIndex);
  }
}
 
void resolveSourceIndicesCollapseShape(Location loc, PatternRewriter &rewriter,
                                       memref::CollapseShapeOp collapseShapeOp,
                                       ValueRange indices,
                                       SmallVectorImpl<Value> &sourceIndices,
                                       bool startsInbounds) {
  // Note: collapse_shape requires a strided memref, we can do this.
  auto metadata = memref::ExtractStridedMetadataOp::create(
      rewriter, loc, collapseShapeOp.getSrc());
  SmallVector<OpFoldResult> sourceSizes = metadata.getConstifiedMixedSizes();
  for (auto [index, group] :
       llvm::zip(indices, collapseShapeOp.getReassociationIndices())) {
    assert(!group.empty() && "association indices groups cannot be empty");
    int64_t groupSize = group.size();
 
    if (groupSize == 1) {
      sourceIndices.push_back(index);
      continue;
    }
 
    // If we don't know that this value is in-bounds, the largest return value
    // of the delinearization may exceed `sourceSizes[d]`, so we drop that first
    // group entry in order to maintain soundness.
    auto trimmedGroup =
        ArrayRef<int64_t>(group).drop_front(startsInbounds ? 0 : 1);
    SmallVector<OpFoldResult> basis = llvm::map_to_vector(
        trimmedGroup, [&](int64_t d) { return sourceSizes[d]; });
    auto delinearize = affine::AffineDelinearizeIndexOp::create(
        rewriter, loc, index, basis, /*hasOuterBound=*/startsInbounds);
    llvm::append_range(sourceIndices, delinearize.getResults());
  }
  if (collapseShapeOp.getReassociationIndices().empty()) {
    auto zeroAffineMap = rewriter.getConstantAffineMap(0);
    int64_t srcRank =
        cast<MemRefType>(collapseShapeOp.getViewSource().getType()).getRank();
    OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
        rewriter, loc, zeroAffineMap, ArrayRef<OpFoldResult>{});
    for (int64_t i = 0; i < srcRank; i++) {
      sourceIndices.push_back(
          getValueOrCreateConstantIndexOp(rewriter, loc, ofr));
    }
  }
}
 
LogicalResult resolveSourceIndicesRankReducingSubview(
    Location loc, OpBuilder &b, memref::SubViewOp subViewOp, ValueRange indices,
    SmallVectorImpl<Value> &sourceIndices) {
  if (!subViewOp.hasZeroOffset() || !subViewOp.hasUnitStride())
    return failure();
 
  MemRefType srcType = subViewOp.getSourceType();
  MemRefType resType = subViewOp.getType();
  unsigned srcRank = srcType.getRank();
  unsigned resRank = resType.getRank();
  if (srcRank <= resRank || indices.size() != resRank)
    return failure();
 
  auto droppedDims = subViewOp.getDroppedDims();
  if (droppedDims.none() || droppedDims.count() != srcRank - resRank)
    return failure();
 
  auto mixedSizes = subViewOp.getMixedSizes();
  if (mixedSizes.size() != srcRank)
    return failure();
 
  unsigned resultDim = 0;
  for (unsigned sourceDim = 0; sourceDim < srcRank; ++sourceDim) {
    if (droppedDims.test(sourceDim)) {
      auto sizeCst = getConstantIntValue(mixedSizes[sourceDim]);
      if (!sizeCst || *sizeCst != 1)
        return failure();
      sourceIndices.push_back(
          getValueOrCreateConstantIndexOp(b, loc, b.getIndexAttr(0)));
      continue;
    }
    if (resultDim >= indices.size())
      return failure();
    sourceIndices.push_back(indices[resultDim++]);
  }
  if (resultDim != indices.size())
    return failure();
 
  return success();
}
 
bool hasNegativeStaticStride(MemRefType memRefTy) {
  auto [strides, offset] = memRefTy.getStridesAndOffset();
  return llvm::any_of(strides, [](int64_t stride) {
    return ShapedType::isStatic(stride) && stride < 0;
  });
}
 
namespace {
// TODO: This footprint modelling and its logic overlap with
// `SubsetOpInterface`, which is unfortunately tensor-only today; if it gains
// memref support, look into unifying with it to improve code reuse.
//
/// Footprint of a memref value relative to a `root` memref reached through a
/// chain of rank-preserving `memref.subview` ops whose `strides` operand is all
/// ones (so each result index steps one source element in every dimension): per
/// `root` dimension, the `[offset, offset + size)` interval the value covers (a
/// `size` may be dynamic). The root is the underlying buffer if the whole chain
/// composes, otherwise the result of the first subview (walking toward the
/// root) that cannot be composed. Two footprints are therefore comparable only
/// when they resolve to the *same* root value.
struct SubviewFootprint {
  Value root;
  SmallVector<int64_t> offsets;
  SmallVector<int64_t> sizes;
};
} // namespace
 
/// Resolve `base` through a `memref.subview` chain to a footprint, or
/// `std::nullopt` if it cannot be modelled. Composition stops at the first
/// subview with a dynamic offset, a `strides` operand other than all ones, or a
/// rank reduction; that subview's own result becomes the footprint root. This
/// keeps two slices that share a common (possibly dynamically-offset) base
/// comparable through their static offsets relative to that base. Sizes may be
/// dynamic: such dimensions simply cannot be used to prove disjointness.
static std::optional<SubviewFootprint> resolveSubviewFootprint(Value base) {
  auto baseTy = dyn_cast<MemRefType>(base.getType());
  if (!baseTy)
    return std::nullopt;
  SubviewFootprint footprint;
  // Size is fixed by `base` itself; offset accumulates the static offsets of
  // the composable subviews walked toward the root.
  footprint.offsets.assign(baseTy.getRank(), 0);
  footprint.sizes.assign(baseTy.getShape().begin(), baseTy.getShape().end());
 
  Value cur = base;
  while (auto sv = cur.getDefiningOp<SubViewOp>()) {
    ArrayRef<int64_t> staticOffsets = sv.getStaticOffsets();
    ArrayRef<int64_t> staticStrides = sv.getStaticStrides();
    // Rank-reducing subviews, or those with a dynamic offset or a `strides`
    // operand other than all ones, cannot be composed: stop here and use `sv`'s
    // result as the root.
    if (sv.getSourceType().getRank() != sv.getType().getRank() ||
        llvm::any_of(staticOffsets, ShapedType::isDynamic) ||
        llvm::any_of(staticStrides, [](int64_t s) { return s != 1; }))
      break;
    for (unsigned d = 0, e = staticOffsets.size(); d < e; ++d)
      footprint.offsets[d] += staticOffsets[d];
    cur = sv.getSource();
  }
  // The root must be a genuine buffer or a `memref.subview` result, not another
  // kind of view (collapse/expand/reshape) whose coordinate remapping would
  // invalidate the per-dimension offset bookkeeping.
  if (auto *def = cur.getDefiningOp())
    if (isa<ViewLikeOpInterface>(def) && !isa<SubViewOp>(def))
      return std::nullopt;
  footprint.root = cur;
  return footprint;
}
 
/// True if two footprints provably cover disjoint memory: they must share the
/// same root and be separated along at least one statically-sized dimension.
static bool areFootprintDisjoint(const SubviewFootprint &a,
                                 const SubviewFootprint &b) {
  if (a.root != b.root)
    return false;
  for (unsigned d = 0, e = a.offsets.size(); d < e; ++d) {
    // A dimension can prove disjointness only when both extents are static.
    if (ShapedType::isDynamic(a.sizes[d]) || ShapedType::isDynamic(b.sizes[d]))
      continue;
    int64_t aHi = a.offsets[d] + a.sizes[d];
    int64_t bHi = b.offsets[d] + b.sizes[d];
    if (aHi <= b.offsets[d] || bHi <= a.offsets[d])
      return true;
  }
  return false;
}
 
bool hasNoAliasingAccessInScope(Value base, Operation *scope,
                                ArrayRef<Operation *> excludedOps,
                                bool readsAreSafe) {
  auto baseMemref = dyn_cast<MemrefValue>(base);
  if (!baseMemref)
    return false;
  Value rootBuffer = skipViewLikeOps(baseMemref);
  // Footprint of `base`, if modellable; enables the disjointness escape.
  std::optional<SubviewFootprint> baseFp = resolveSubviewFootprint(base);
 
  // Visit the transitive users of the root buffer, following views.
  SmallVector<Operation *> worklist(rootBuffer.getUsers().begin(),
                                    rootBuffer.getUsers().end());
  SmallPtrSet<Operation *, 8> processed;
  while (!worklist.empty()) {
    Operation *user = worklist.pop_back_val();
    if (!processed.insert(user).second)
      continue;
    // Views do not access memory; follow them to their own users.
    if (auto viewLike = dyn_cast<ViewLikeOpInterface>(user)) {
      Value viewDest = viewLike.getViewDest();
      worklist.append(viewDest.getUsers().begin(), viewDest.getUsers().end());
      continue;
    }
    if (isMemoryEffectFree(user) || llvm::is_contained(excludedOps, user))
      continue;
    if (!scope->isAncestor(user))
      continue;
    // In-`scope`, memory-effecting op: each of its operands that resolves to
    // the root buffer conflicts unless it is a provably-disjoint slice or (when
    // allowed) is only read.
    for (Value operand : user->getOperands()) {
      auto slice = dyn_cast<MemrefValue>(operand);
      if (!slice || skipViewLikeOps(slice) != rootBuffer)
        continue;
      if (baseFp) {
        std::optional<SubviewFootprint> sliceFp =
            resolveSubviewFootprint(slice);
        if (sliceFp && areFootprintDisjoint(*baseFp, *sliceFp))
          continue;
      }
      if (readsAreSafe && isa<MemoryEffectOpInterface>(user) &&
          !hasEffect<MemoryEffects::Write>(user, slice))
        continue;
      return false;
    }
  }
  return true;
}
 
} // namespace memref
} // namespace mlir