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

[Matthias Springer]

================
@@ -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(
----------------
matthias-springer wrote:

I thought the `AffineMapAttr` is semi-deprecated. We use `StridedLayoutAttr` in most places.

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