[mlir][MemRef] Add early exit for computing dropped unit-dims.

[mravishankar]

Summary

Computing dropped unit-dims when all the unit dims are dropped, does
not need to check for strides being dropped.
This also enables canonicalization of reduced-rank subviews.

[mravishankar]

Rebase

[NathanielMcVicar]

I believe this change produces a comparison warning with clang:

llvm-project/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp:764:50: error: comparison of integers of different signs: 'unsigned long' and 'int64_t' (aka 'long') [-Werror,-Wsign-compare]
  if (unusedDims.count() + reducedType.getRank() == originalType.getRank())

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType,
   if (originalType.getRank() == reducedType.getRank())
     return unusedDims;
 
   for (const auto &dim : llvm::enumerate(sizes))
     if (auto attr = dim.value().dyn_cast<Attribute>())
       if (attr.cast<IntegerAttr>().getInt() == 1)
         unusedDims.set(dim.index());
 
+  // Early exit for the case where the number of unused dims matches the number
+  // of ranks reduced.
+  if (unusedDims.count() + reducedType.getRank() == originalType.getRank())
+    return unusedDims;
+
   SmallVector<int64_t> originalStrides, candidateStrides;
   int64_t originalOffset, candidateOffset;
   if (failed(
           getStridesAndOffset(originalType, originalStrides, originalOffset)) ||
       failed(
           getStridesAndOffset(reducedType, candidateStrides, candidateOffset)))
     return llvm::None;
 

mlir/test/Dialect/MemRef/canonicalize.mlir

 //  CHECK-SAME: (%[[ARG:.*]]: memref<?xi8>, %[[SIZE1:.*]]: index, %[[SIZE2:.*]]: index)
 //       CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [%[[SIZE2]]], strides: [1]
 //       CHECK: return %[[RES]]
 func @reinterpret_of_subview(%arg : memref<?xi8>, %size1: index, %size2: index) -> memref<?xi8> {
   %0 = memref.subview %arg[0] [%size1] [1] : memref<?xi8> to memref<?xi8>
   %1 = memref.reinterpret_cast %0 to offset: [0], sizes: [%size2], strides: [1] : memref<?xi8> to memref<?xi8>
   return %1 : memref<?xi8>
 }
+
+// -----
+
+func @canonicalize_rank_reduced_subview(%arg0 : memref<8x?xf32>,
+    %arg1 : index) -> memref<?xf32, offset : ?, strides : [?]> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = memref.subview %arg0[%c0, %c0] [1, %arg1] [%c1, %c1] : memref<8x?xf32> to memref<?xf32, offset : ?, strides : [?]>
+  return %0 :  memref<?xf32, offset : ?, strides : [?]>
+}
+//  CHECK-DAG: #[[MAP:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
+//      CHECK: func @canonicalize_rank_reduced_subview
+// CHECK-SAME:     %[[ARG0:.+]]: memref<8x?xf32>
+// CHECK-SAME:     %[[ARG1:.+]]: index
+//      CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[ARG0]][0, 0] [1, %[[ARG1]]] [1, 1]
+// CHECK-SAME:       memref<8x?xf32> to memref<?xf32, #[[MAP]]>

mlir/test/Dialect/MemRef/invalid.mlir

 func @invalid_rank_reducing_subview(%arg0 : memref<?x?xf32>, %arg1 : index, %arg2 : index) {
   // expected-error@+1 {{expected result type to be 'memref<?x1xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>>' or a rank-reduced version. (mismatch of result layout)}}
   %0 = memref.subview %arg0[0, %arg1][%arg2, 1][1, 1] : memref<?x?xf32> to memref<?xf32>
   return
 }
 
 // -----
 
-func @static_stride_to_dynamic_stride(%arg0 : memref<?x?x?xf32>, %arg1 : index,
-    %arg2 : index) -> memref<?x?xf32, offset:?, strides: [?, ?]> {
-  // expected-error @+1 {{expected result type to be 'memref<1x?x?xf32, affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>>' or a rank-reduced version. (mismatch of result layout)}}
-  %0 = memref.subview %arg0[0, 0, 0] [1, %arg1, %arg2] [1, 1, 1] : memref<?x?x?xf32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-  return %0 : memref<?x?xf32, offset: ?, strides: [?, ?]>
-}
-
-// -----
-
 #map0 = affine_map<(d0, d1)[s0] -> (d0 * 16 + d1)>
 
 func @subview_bad_offset_1(%arg0: memref<16x16xf32>) {
   %c0 = arith.constant 0 : index
   %c8 = arith.constant 8 : index
   // expected-error @+1 {{expected result type to be 'memref<8x8xf32, affine_map<(d0, d1)[s0] -> (d0 * 16 + s0 + d1)>>' or a rank-reduced version}}
   %s2 = memref.subview %arg0[%c8, %c8][8, 8][1, 1]  : memref<16x16xf32> to memref<8x8xf32, #map0>
   return