[Mlir-commits] [mlir] [mlir] Use identity map to construct the memref type as possible (PR #183051)

[donald chen]

================
@@ -566,13 +566,46 @@ unsigned MemRefType::getMemorySpaceAsInt() const {
   return detail::getMemorySpaceAsInt(getMemorySpace());
 }
 
+static bool getNumContiguousTrailingDimsImpl(ArrayRef<int64_t> shape,
+                                             ArrayRef<int64_t> strides) {
+  const int64_t n = shape.size();
+  // A memref with dimensions `d0, d1, ..., dn-1` and strides
+  // `s0, s1, ..., sn-1` is contiguous up to dimension `k`
+  // if each stride `si` is the product of the dimensions `di+1, ..., dn-1`,
+  // for `i` in `[k, n-1]`.
+  // Ignore stride elements if the corresponding dimension is 1, as they are
+  // of no consequence.
+  int64_t dimProduct = 1;
+  for (int64_t i = n - 1; i >= 0; --i) {
+    if (shape[i] == 1)
+      continue;
+    if (strides[i] != dimProduct)
+      return n - i - 1;
+    if (shape[i] == ShapedType::kDynamic)
+      return n - i;
+    dimProduct *= shape[i];
+  }
+  return n;
+}
+
 MemRefType MemRefType::get(ArrayRef<int64_t> shape, Type elementType,
                            MemRefLayoutAttrInterface layout,
                            Attribute memorySpace) {
-  // Use default layout for empty attribute.
-  if (!layout)
+  if (!layout) {
+    // Use default layout for empty attribute.
     layout = AffineMapAttr::get(AffineMap::getMultiDimIdentityMap(
         shape.size(), elementType.getContext()));
+  } else {
+    // If the layout can be inferred to be an identity, prefer using the
+    // identity layout.
+    int64_t offset;
+    SmallVector<int64_t> strides;
+    (void)layout.getStridesAndOffset(shape, strides, offset);
+    if (offset == 0 &&
+        getNumContiguousTrailingDimsImpl(shape, strides) == shape.size())
+      layout = AffineMapAttr::get(AffineMap::getMultiDimIdentityMap(
----------------
cxy-1993 wrote:

StridedLayoutAttr carries less information than AffineMapAttr: the latter can identify a memref as 'identity' directly from its type. An identity map enables more optimizations within MLIR, and the motivation behind this patch was specifically to use AffineMapAttr more frequently to leverage those identity-related optimizations. I don't quite understand the reasoning for deprecating AffineMapAttr.

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