[Mlir-commits] [mlir] [memref] Simplify loads from reinterpret_cast of 1D contiguous memrefs (PR #188459)

[Andrzej Warzyński]

================
@@ -196,6 +198,259 @@ struct CopyToScalarLoadAndStore : public OpRewritePattern<memref::CopyOp> {
   }
 };
 
+/// Captures info about MemRefs that are effectively 1D (the leading or trailing
+/// dims are all 1). The only accepted non-unit dim is either the leading of the
+/// trailing dim.
+///
+/// Examples:
+/// memref<1x1x4xf32>, memref<4x1x1xf32>, memref<1x1x1xf32>
+///
+struct ShapeInfoFor1DMemRef {
+  // Are all dims == 1? `false` means that there is exactly one dim != 1.
+  bool allOnes = true;
+  // If there is a non-unit boundary dim, is it the leading or the trailing dim?
+  bool isLeadingDimNonUnit = false;
+};
+
+/// Returns information about a MemRef if it contains at most one non-unit
+/// dimension.
+///
+/// The single non-unit dimension, if present, must be on the left or right
+/// boundary. Rank-1 non-unit MemRefs are treated as being on both boundaries.
+static std::optional<ShapeInfoFor1DMemRef>
+getShapeInfoFor1DMemRef(MemRefType type) {
+  ArrayRef<int64_t> shape = type.getShape();
+  int64_t nonUnitCount =
+      llvm::count_if(shape, [](int64_t dim) { return dim != 1; });
+  // Return default values if missing nonUnitDim
+  if (nonUnitCount == 0)
+    return ShapeInfoFor1DMemRef{};
+  // Return no info if MemRef breaks nonUnitDim requirements (more nonUnitDims)
+  if (nonUnitCount > 1)
+    return std::nullopt;
+  // Return no info if MemRef breaks nonUnitDim requirements (nonUnitDim in
+  // non-boundary pos)
+  if (shape.front() == 1 && shape.back() == 1)
+    return std::nullopt;
+
+  return ShapeInfoFor1DMemRef{/*allOnes=*/false,
+                              /*isLeadingDimNonUnit=*/shape.front() != 1};
+}
+
+static bool hasStaticZeroOffset(memref::ReinterpretCastOp rc) {
+  ArrayRef<int64_t> offsets = rc.getStaticOffsets();
+  // FIXME: Despite what `getStaticOffsets` implies, `reinterpret_cast` takes
+  // only a single offset. That should be fixed at the op definition level.
+  assert(offsets.size() == 1 && "Expecting single offset");
+  return !ShapedType::isDynamic(offsets[0]) && offsets[0] == 0;
+}
+
+static std::optional<int64_t> getConstantIndex(Value v) {
+  if (auto cst = v.getDefiningOp<arith::ConstantIndexOp>())
+    return cst.value();
+  // Non-constant and dynamic indices
+  return std::nullopt;
+}
+
+static bool isConstantIndexExplicitlyOutOfBounds(Value idx,
+                                                 int64_t upperBound) {
+  // Only statically known `arith.constant` indices are checked here.
+  std::optional<int64_t> idxVal = getConstantIndex(idx);
+  return idxVal && (*idxVal < 0 || *idxVal >= upperBound);
+}
+
+/// Examples accepted by this shape restriction:
+///   memref<999xf32>       <-> memref<1x1x999xf32>
+///   memref<1x108xf32>     <-> memref<1x1x1x108xf32>
+///   memref<100x1xf32>     <-> memref<100x1x1xf32>
+///   memref<1>             <-> memref<1x1x1>
+///
+/// General reinterpret_casts are intentionally rejected.
+static bool isPureRankExpansionOrCollapsingRC(memref::ReinterpretCastOp rc) {
+  auto inputTy = cast<MemRefType>(rc.getSource().getType());
+  auto outputTy = cast<MemRefType>(rc.getResult().getType());
+
+  // Only zero, statically known offsets are accepted. Non-zero or dynamic
+  // offsets would require reasoning about storage shifts in the underlying
+  // reinterpret_cast, which this helper does not model.
+  if (!hasStaticZeroOffset(rc))
+    return false;
+
+  // Dynamic sizes/strides prevent precise reasoning about the underlying
+  // reinterpret_cast, so only fully static shape metadata is accepted.
+  if (llvm::any_of(rc.getStaticSizes(), ShapedType::isDynamic) ||
+      llvm::any_of(rc.getStaticStrides(), ShapedType::isDynamic))
+    return false;
+
+  // Only shapes with at most one non-unit dimension are accepted. This rules
+  // out more general multi-dimensional reinterpret_casts and restricts the
+  // helper to unit-dim insertion/removal around a single logical dimension.
+  std::optional<ShapeInfoFor1DMemRef> inputNonUnitDim =
+      getShapeInfoFor1DMemRef(inputTy);
+  std::optional<ShapeInfoFor1DMemRef> outputNonUnitDim =
+      getShapeInfoFor1DMemRef(outputTy);
+  // Bail out if either type does not satisfy the single-boundary-non-unit-dim
+  // restriction described above.
+  if (!inputNonUnitDim || !outputNonUnitDim)
+    return false;
+
+  // The source and result must either both have a single non-unit dimension
+  // or both be all-ones.
+  if (inputNonUnitDim->allOnes != outputNonUnitDim->allOnes)
+    return false;
+  if (inputNonUnitDim->allOnes)
+    return true;
+
+  // The preserved non-unit dimension must have the same size.
+  if (inputTy.getDimSize(
+          inputNonUnitDim->isLeadingDimNonUnit ? 0 : inputTy.getRank() - 1) !=
+      outputTy.getDimSize(
+          outputNonUnitDim->isLeadingDimNonUnit ? 0 : outputTy.getRank() - 1))
+    return false;
+
+  // If both sides have rank > 1, the non-unit dimension must be on the same
+  // boundary. Rank-1 MemRefs are accepted against either boundary.
+  if (inputTy.getRank() != 1 && outputTy.getRank() != 1 &&
+      inputNonUnitDim->isLeadingDimNonUnit !=
+          outputNonUnitDim->isLeadingDimNonUnit)
+    return false;
+
+  return true;
+}
+
+/// Checks statically known indices accessed by a load from a pure rank
+/// expansion/collapsing to ensure in-bounds only access. Dynamic indices are
+/// accepted.
----------------
banach-space wrote:

```suggestion
/// Checks statically known and constant indices accessed by a load from a pure rank
/// expansion/collapsing to ensure in-bounds only access. Fully dynamic indices are
/// skipped (there is no way to verify them).
```

1. "accepted" alone is a bit confusing - "accepted" as what?
2. From the point of view of this check, there are three types of indices: static (e.g. `12`), dynamic ssa values that represent constants (e.g. `%idx = arith.constant 0`) and fully dynamic ssa values (e.g. function arguments). It would be good to differentiate.

https://github.com/llvm/llvm-project/pull/188459