[fixup] Misc NFC changes

 - rename `getMaxCollapsabelTrailingDims` to `getMaxContiguousTrailingDims`
 - new set of examples
 - remove redundant call to `isIdentify()`
 - make sure a variable type is visible on the declaration line
 - some micro-optimisation

Author: momchil-velikov

mlir/include/mlir/IR/BuiltinTypes.td

@@ -838,19 +838,25 @@ def Builtin_MemRef : Builtin_Type<"MemRef", "memref", [
     ///
     bool areTrailingDimsContiguous(int64_t n);
 
-    /// Return the maximum number of trailing dimensions that can be
-    /// collapsed.
+    /// Return the maximum number of trailing dimensions that are
+    /// contiguous.
     ///
     /// Examples:
-    ///   - memref<2x3x2xi8, strided<[24, 12, 2]>, the number of collapsable
-    ///     trailing dimensions is 0
-    ///   - memref<2x3x2xi8, strided<[12, 6, 1]>, the number of collapsable
+    ///   - memref<5x3x2xi8, strided<[6,2,1]>>, the number of collapsable
     ///     trailing dimensions is 3
-    ///   - memref<5x4x3x2xi8, strided<[48, 6, 2, 1]>, the number of
-    ///     collapsable trailing dimensions is 2.
-    ///   - memref<5x4x?x2xi8>, the number of collapsable trailing dimensions
-    ///     is 4.
-    int64_t getMaxCollapsableTrailingDims();
+    ///   - memref<5x3x2xi8, strided<[12,2,1]>>, the number of collapsable
+    ///     trailing dimensions is 2 (dimension 0 is non-contiguous)
+    ///   - memref<5x3x2xi8, strided<[12,4,1]>>, the number of collapsable
+    ///     trailing dimensions is 1 (dimension 1 is non-contiguous)
+    ///   - memref<5x3x2xi8, strided<[12,4,2]>>, the number of collapsable
+    ///     trailing dimensions is 0 (dimension 2 is non-contiguous)
+    ///   - memref<?x3x2xi8, strided<[6,2,1]>>, the number of collapsable
+    ///     trailing dimensions is 3
+    ///   - memref<?x3x2xi8, strided<[12,2,1]>>, the number of collapsable
+    ///     trailing dimensions is 2 (dimension 0 is non-contiguous)
+    ///   - memref<5x?x2xi8, strided<[?,2,1]>>, the number of collapsable
+    ///     trailing dimensions is 2 (stride 0 is dynamic)
+    int64_t getMaxContiguousTrailingDims();
 
     /// Return a version of this type with identity layout if it can be
     /// determined statically that the layout is the canonical contiguous

mlir/lib/IR/BuiltinTypes.cpp

@@ -646,11 +646,10 @@ LogicalResult MemRefType::verify(function_ref<InFlightDiagnostic()> emitError,
 }
 
 bool MemRefType::areTrailingDimsContiguous(int64_t n) {
-  return getLayout().isIdentity() ||
-         getMaxCollapsableTrailingDims() >= std::min(n, getRank());
+  return getMaxContiguousTrailingDims() >= std::min(n, getRank());
 }
 
-int64_t MemRefType::getMaxCollapsableTrailingDims() {
+int64_t MemRefType::getMaxContiguousTrailingDims() {
   const int64_t n = getRank();
 
   // memrefs with identity layout are entirely contiguous.
@@ -664,7 +663,7 @@ int64_t MemRefType::getMaxCollapsableTrailingDims() {
   if (!succeeded(getStridesAndOffset(strides, offset)))
     return 0;
 
-  auto shape = getShape();
+  ArrayRef<int64_t> shape = getShape();
 
   // A memref with dimensions `d0, d1, ..., dn-1` and strides
   // `s0, s1, ..., sn-1` is contiguous up to dimension `k`
@@ -674,6 +673,8 @@ int64_t MemRefType::getMaxCollapsableTrailingDims() {
   for (int64_t i = n - 1; i >= 0; --i) {
     if (strides[i] != dimProduct)
       return n - i - 1;
+    if (shape[i] == 1)
+      continue;
     if (shape[i] == ShapedType::kDynamic)
       return n - i;
     dimProduct *= shape[i];

mlir/unittests/Dialect/MemRef/LayoutTest.cpp

@@ -27,55 +27,55 @@ TEST(MemRefLayout, maxCollapseDim) {
 
   // memref<2x2x2xf32, strided<[4,2,1]>
   auto m1 = MemRefType::get({2, 2, 2}, f32, strided({4, 2, 1}));
-  EXPECT_EQ(m1.getMaxCollapsableTrailingDims(), 3);
+  EXPECT_EQ(m1.getMaxContiguousTrailingDims(), 3);
 
   // memref<2x2x2xf32, strided<[8,2,1]>
   auto m2 = MemRefType::get({2, 2, 2}, f32, strided({8, 2, 1}));
-  EXPECT_EQ(m2.getMaxCollapsableTrailingDims(), 2);
+  EXPECT_EQ(m2.getMaxContiguousTrailingDims(), 2);
 
   // memref<2x2x2xf32, strided<[8,4,1]>
   auto m3 = MemRefType::get({2, 2, 2}, f32, strided({8, 4, 1}));
-  EXPECT_EQ(m3.getMaxCollapsableTrailingDims(), 1);
+  EXPECT_EQ(m3.getMaxContiguousTrailingDims(), 1);
 
   // memref<2x2x2xf32, strided<[8,4,2]>
   auto m4 = MemRefType::get({2, 2, 2}, f32, strided({8, 4, 2}));
-  EXPECT_EQ(m4.getMaxCollapsableTrailingDims(), 0);
+  EXPECT_EQ(m4.getMaxContiguousTrailingDims(), 0);
 
   // memref<2x2x?xf32, strided<[?,?,1]>
   auto m5 = MemRefType::get({2, 2, _}, f32, strided({_, _, 1}));
-  EXPECT_EQ(m5.getMaxCollapsableTrailingDims(), 1);
+  EXPECT_EQ(m5.getMaxContiguousTrailingDims(), 1);
 
   // memref<2x2x?xf32, strided<[?,?,2]>
   auto m6 = MemRefType::get({2, 2, _}, f32, strided({_, _, 2}));
-  EXPECT_EQ(m6.getMaxCollapsableTrailingDims(), 0);
+  EXPECT_EQ(m6.getMaxContiguousTrailingDims(), 0);
 
   // memref<2x?x2xf32, strided<[?,2,1]>
   auto m7 = MemRefType::get({2, _, 2}, f32, strided({_, 2, 1}));
-  EXPECT_EQ(m7.getMaxCollapsableTrailingDims(), 2);
+  EXPECT_EQ(m7.getMaxContiguousTrailingDims(), 2);
 
   // memref<2x?x2xf32, strided<[?,4,1]>
   auto m8 = MemRefType::get({2, _, 2}, f32, strided({_, 4, 1}));
-  EXPECT_EQ(m8.getMaxCollapsableTrailingDims(), 1);
+  EXPECT_EQ(m8.getMaxContiguousTrailingDims(), 1);
 
   // memref<2x?x2xf32, strided<[?,4,2]>
   auto m9 = MemRefType::get({2, _, 2}, f32, strided({_, 4, 2}));
-  EXPECT_EQ(m9.getMaxCollapsableTrailingDims(), 0);
+  EXPECT_EQ(m9.getMaxContiguousTrailingDims(), 0);
 
   // memref<?x2x2xf32, strided<[4,2,1]>
   auto m10 = MemRefType::get({_, 2, 2}, f32, strided({4, 2, 1}));
-  EXPECT_EQ(m10.getMaxCollapsableTrailingDims(), 3);
+  EXPECT_EQ(m10.getMaxContiguousTrailingDims(), 3);
 
   // memref<?x2x2xf32, strided<[8,2,1]>
   auto m11 = MemRefType::get({_, 2, 2}, f32, strided({8, 2, 1}));
-  EXPECT_EQ(m11.getMaxCollapsableTrailingDims(), 2);
+  EXPECT_EQ(m11.getMaxContiguousTrailingDims(), 2);
 
   // memref<?x2x2xf32, strided<[8,4,1]>
   auto m12 = MemRefType::get({_, 2, 2}, f32, strided({8, 4, 1}));
-  EXPECT_EQ(m12.getMaxCollapsableTrailingDims(), 1);
+  EXPECT_EQ(m12.getMaxContiguousTrailingDims(), 1);
 
   // memref<?x2x2xf32, strided<[8,4,2]>
   auto m13 = MemRefType::get({_, 2, 2}, f32, strided({8, 4, 2}));
-  EXPECT_EQ(m13.getMaxCollapsableTrailingDims(), 0);
+  EXPECT_EQ(m13.getMaxContiguousTrailingDims(), 0);
 }
 
 TEST(MemRefLayout, contigTrailingDim) {
@@ -176,14 +176,14 @@ TEST(MemRefLayout, identityMaps) {
 
   // memref<2x2x2xf32>
   auto m1 = MemRefType::get({2, 2, 2}, f32);
-  EXPECT_EQ(m1.getMaxCollapsableTrailingDims(), 3);
+  EXPECT_EQ(m1.getMaxContiguousTrailingDims(), 3);
   EXPECT_TRUE(m1.areTrailingDimsContiguous(1));
   EXPECT_TRUE(m1.areTrailingDimsContiguous(2));
   EXPECT_TRUE(m1.areTrailingDimsContiguous(3));
 
   // memref<?x?x?xf32>
   auto m2 = MemRefType::get({_, _, _}, f32);
-  EXPECT_EQ(m2.getMaxCollapsableTrailingDims(), 3);
+  EXPECT_EQ(m2.getMaxContiguousTrailingDims(), 3);
   EXPECT_TRUE(m2.areTrailingDimsContiguous(1));
   EXPECT_TRUE(m2.areTrailingDimsContiguous(2));
   EXPECT_TRUE(m2.areTrailingDimsContiguous(3));