[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)

[Andrzej Warzyński]

================

----------------
banach-space wrote:

```suggestion
  bool isScalarCopy = !llvm::all_of(resType.getShape(),
                    [](int64_t size) { return size == 1; }));
  
  // For scalar copies, result strides are irrelevant, including dynamic ones.
  // For non-scalar copies, require static result strides identical to the
  // identity strides of the reinterpret_cast source.
  if (isScalarCopy) {
    SmallVector<int64_t> srcIdentityStrides =
        computeStrides(srcType.getShape());
    ArrayRef<int64_t> rcResultStrides = rc.getStaticStrides();

    assert((srcIdentityStrides.size() == rcResultStrides.size()) &&
           "Expecting same number of strides for rank-preserving "
           "reinterpret_casts.");

    if (!llvm::all_of(llvm::zip_equal(srcIdentityStrides, rcResultStrides),
                      [](auto pair) {
                        auto [srcStride, resultStride] = pair;
                        return !ShapedType::isDynamic(resultStride) &&
                               srcStride == resultStride;
}))
      return std::nullopt;
      }
   
   // TODO - Add comment explaining _why_ non-unit sizes are tracked in this case   
   if (isScalarCopy) {
    for (auto [dim, resultSize] : llvm::enumerate(resType.getShape())) {
      if (resultSize != 1)
        dimsAndOffs.nonUnitDimsPos.push_back(static_cast<unsigned>(dim));
    }
  }
```

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