[Mlir-commits] [mlir] [mlir][memref] Rewrite scalar `memref.copy` through reinterpret_cast into load/store (PR #186118)

[Andrzej Warzyński]

================
@@ -0,0 +1,233 @@
+//===-ElideReinterpretCast.cpp - Expansion patterns for MemRef operations-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include <cassert>
+
+namespace mlir {
+namespace memref {
+#define GEN_PASS_DEF_ELIDEREINTERPRETCASTPASS
+#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
+} // namespace memref
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+/// Returns true if `rc` represents a scalar view (all sizes == 1)
+/// into a memref that has exactly one non-unit dimension located at
+/// either the first or last position (i.e. a "row" or "column").
+///
+/// Examples that return true:
+///
+///   // Row-major slice (last dim is non-unit)
+///   memref.reinterpret_cast %buff to offset: [%off],
+///     sizes: [1, 1, 1], strides: [1, 1, 1]
+///     : memref<1x1x8xi32> to memref<1x1x1xi32>
+///
+///   // Column-major slice (first dim is non-unit)
+///   memref.reinterpret_cast %buff to offset: [%off],
+///     sizes: [1, 1], strides: [1, 1]
+///     : memref<2x1xf32> to memref<1x1xf32>
+///
+///   // Random strides
+///   memref.reinterpret_cast %buff to offset: [%off],
+///     sizes: [1, 1], strides: [10, 100]
+///     : memref<2x1xf32, strided<[10, 100]>>
+///         to memref<1x1xf32>
+///
+///   // Rank-1 case
+///   memref.reinterpret_cast %buf to offset: [%off],
+///     sizes: [1], strides: [1]
+///     : memref<8xi32> to memref<1xi32>
+///
+/// Examples that return false:
+///
+///   // More non-unit dims
+///   memref.reinterpret_cast %buff to offset: [%off],
+///     sizes: [1, 1, 1], strides: [1, 1, 1]
+///     : memref<1x2x8xi32> to memref<1x1x1xi32>
+///
+///   // View is not scalar (size != 1)
+///   memref.reinterpret_cast %buff to offset: [%off],
+///     sizes: [2, 1], strides: [1, 1]
+///     : memref<1x2xf32> to memref<2x1xf32>
+///
+///   // Base has non-identity layout
+///   %buff = memref.alloc() : memref<1x2xf32, strided<[1, 3]>>
+///   memref.reinterpret_cast %buff to offset: [%off],
+///     sizes: [1, 1], strides: [1, 1]
+///     : memref<1x2xf32, strided<[1, 3]>> to memref<1x1xf32>
----------------
banach-space wrote:

IIUC, some of the restrictions could be lifted? Lets document that. Otherwise it might sound as if that would be very hard and/or impossible.

```suggestion
///     : memref<1x2xf32, strided<[1, 3]>> to memref<1x1xf32>
///
/// TODO: Relax some of the conditions above.
```

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