//===- MemRefUtils.h - MemRef transformation utilities ----------*- C++ -*-===//
//
// 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 header file defines prototypes for various transformation utilities for
// the MemRefOps dialect. These are not passes by themselves but are used
// either by passes, optimization sequences, or in turn by other transformation
// utilities.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H
#define MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H

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

namespace mlir {

class MemRefType;

/// A value with a memref type.
using MemrefValue = TypedValue<BaseMemRefType>;

namespace memref {

/// Returns true, if the memref type has static shapes and represents a
/// contiguous chunk of memory.
bool isStaticShapeAndContiguousRowMajor(MemRefType type);

/// Controls how the per-dimension contribution to `linearizedSize` is divided
/// by `dstBits / srcBits` when scaling down to the emulated type. The offset
/// and intra-data offset are unaffected; they always use floor division and
/// remainder respectively.
/// - `Floor`: round each `stride * size / scaler` down. Suitable for indexing
///   computations where a partial trailing byte is not included.
/// - `Ceil`: round up, matching the result-shape size used by narrow-type
///   memref type conversion (see `getLinearizedShape`). Use this when the
///   caller needs the linearized size to cover all source elements, e.g. when
///   building the size attribute of a converted `memref.reinterpret_cast`.
enum class LinearizedDivKind { Floor, Ceil };

/// For a `memref` with `offset`, `sizes` and `strides`, returns the
/// offset, size, and potentially the size padded at the front to use for the
/// linearized `memref`.
/// - If the linearization is done for emulating load/stores of
///   element type with bitwidth `srcBits` using element type with
///   bitwidth `dstBits`, the linearized offset and size are
///   scaled down by `dstBits`/`srcBits`.
/// - If `indices` is provided, it represents the position in the
///   original `memref` being accessed. The method then returns the
///   index to use in the linearized `memref`. The linearized index
///   is also scaled down by `dstBits`/`srcBits`. If `indices` is not provided
///   0, is returned for the linearized index.
/// - If the size of the load/store is smaller than the linearized memref
///   load/store, the memory region emulated is larger than the actual memory
///   region needed. `intraDataOffset` returns the element offset of the data
///   relevant at the beginning.
/// - `sizeDivKind` selects floor vs ceil rounding for the `linearizedSize`
///   contribution from each dimension (see `LinearizedDivKind`).
struct LinearizedMemRefInfo {
  OpFoldResult linearizedOffset;
  OpFoldResult linearizedSize;
  OpFoldResult intraDataOffset;
};
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 = LinearizedDivKind::Floor);

/// For a `memref` with `offset` and `sizes`, returns the
/// offset and size to use for the linearized `memref`, assuming that
/// the strides are computed from a row-major ordering of the sizes;
/// - If the linearization is done for emulating load/stores of
///   element type with bitwidth `srcBits` using element type with
///   bitwidth `dstBits`, the linearized offset and size are
///   scaled down by `dstBits`/`srcBits`.
/// - `sizeDivKind` selects floor vs ceil rounding for the `linearizedSize`
///   contribution from each dimension (see `LinearizedDivKind`).
LinearizedMemRefInfo getLinearizedMemRefOffsetAndSize(
    OpBuilder &builder, Location loc, int srcBits, int dstBits,
    OpFoldResult offset, ArrayRef<OpFoldResult> sizes,
    LinearizedDivKind sizeDivKind = LinearizedDivKind::Floor);

/// Track temporary allocations that are never read from. If this is the case
/// it means both the allocations and associated stores can be removed.
void eraseDeadAllocAndStores(RewriterBase &rewriter, Operation *parentOp);

/// Given a set of sizes, return the suffix product.
///
/// When applied to slicing, this is the calculation needed to derive the
/// strides (i.e. the number of linear indices to skip along the (k-1) most
/// minor dimensions to get the next k-slice).
///
/// This is the basis to linearize an n-D offset confined to `[0 ... sizes]`.
///
/// Assuming `sizes` is `[s0, .. sn]`, return the vector<Value>
///   `[s1 * ... * sn, s2 * ... * sn, ..., sn, 1]`.
///
/// It is the caller's responsibility to provide valid OpFoldResult type values
/// and construct valid IR in the end.
///
/// `sizes` elements are asserted to be non-negative.
///
/// Return an empty vector if `sizes` is empty.
///
/// The function emits an IR block which computes suffix product for provided
/// sizes.
SmallVector<OpFoldResult>
computeSuffixProductIRBlock(Location loc, OpBuilder &builder,
                            ArrayRef<OpFoldResult> sizes);
inline SmallVector<OpFoldResult>
computeStridesIRBlock(Location loc, OpBuilder &builder,
                      ArrayRef<OpFoldResult> sizes) {
  return computeSuffixProductIRBlock(loc, builder, sizes);
}

/// Walk up the source chain until an operation that changes/defines the view of
/// memory is found (i.e. skip operations that alias the entire view).
MemrefValue skipFullyAliasingOperations(MemrefValue source);

/// Checks if two (memref) values are the same or statically known to alias
/// the same region of memory.
inline bool isSameViewOrTrivialAlias(MemrefValue a, MemrefValue b) {
  return skipFullyAliasingOperations(a) == skipFullyAliasingOperations(b);
}

/// Walk up the source chain until we find an operation that is not a view of
/// the source memref (i.e. implements ViewLikeOpInterface).
MemrefValue skipViewLikeOps(MemrefValue source);

/// Given the 'indices' of a load/store operation where the memref is a result
/// of a expand_shape op, returns the indices w.r.t to the source memref of the
/// expand_shape op into `sourceIndices`. For example
///
/// %0 = ... : memref<12x42xf32>
/// %1 = memref.expand_shape %0 [[0, 1], [2]]
///    : memref<12x42xf32> into memref<2x6x42xf32>
/// %2 = load %1[%i1, %i2, %i3] : memref<2x6x42xf32
///
/// could be folded into
///
/// %2 = load %0[6 * i1 + i2, %i3] :
///          memref<12x42xf32>
///
/// If `startsInbounds` is true, optimizations that rely on all indices being
/// non-negative and less than the corresponding memref dimension may be
/// performed.
void resolveSourceIndicesExpandShape(Location loc, PatternRewriter &rewriter,
                                     memref::ExpandShapeOp expandShapeOp,
                                     ValueRange indices,
                                     SmallVectorImpl<Value> &sourceIndices,
                                     bool startsInbounds);

/// Given the 'indices' of a load/store operation where the memref is a result
/// of a collapse_shape op, returns the indices w.r.t to the source memref of
/// the collapse_shape op, returing them into `sourceIndices`. For example
///
/// %0 = ... : memref<2x6x42xf32>
/// %1 = memref.collapse_shape %0 [[0, 1], [2]]
///    : memref<2x6x42xf32> into memref<12x42xf32>
/// %2 = load %1[%i1, %i2] : memref<12x42xf32>
///
/// could be folded into
///
/// %2 = load %0[%i1 / 6, %i1 % 6, %i2] :
///          memref<2x6x42xf32>
///
/// If `startsInbounds` is true, optimizations that rely on all indices being
/// non-negative and less than the corresponding memref dimension may be
/// performed.
void resolveSourceIndicesCollapseShape(Location loc, PatternRewriter &rewriter,
                                       memref::CollapseShapeOp collapseShapeOp,
                                       ValueRange indices,
                                       SmallVectorImpl<Value> &sourceIndices,
                                       bool startsInbounds);

/// Given the 'indices' of a load/store operation where the memref is a result
/// of a rank-reducing full subview op, returns the indices w.r.t to the source
/// memref of the memref.subview op. For example
///
///  %alias = memref.subview %src[0, 0, 0][1, 2, 2][1, 1, 1]: memref<1x2x2xf32>
///                           to memref<2x2xf32>
///  %val = memref.load %alias[%i, %j] : memref<2x2xf32>
///
/// could be folded into
///
///  %val = memref.load %src[0, %i, %j] : memref<1x2x2xf32>
LogicalResult resolveSourceIndicesRankReducingSubview(
    Location loc, OpBuilder &b, memref::SubViewOp subViewOp, ValueRange indices,
    SmallVectorImpl<Value> &sourceIndices);

} // namespace memref
} // namespace mlir

#endif // MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H