[mlir][memref] Tighten verification of memref.reinterpret_cast

[herhut]

Summary

We allow the omission of a map in memref.reinterpret_cast under the assumption, that the cast might cast to an identity layout. This change adds verification that the static knowledge that is present in the reinterpret_cast supports this assumption.

[nicolasvasilache]

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

1166

Can you just replace the condition by: if (isStrided(resultType)) ?

[herhut]

Simplified as suggested

[herhut]

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

1166

Good point, I combined the two cases. I was not aware that getStridesAndOffset also works for identity maps. But of course it needs to.

[nicolasvasilache]

Nice cleanup!

mlir/test/mlir-cpu-runner/copy.mlir

38–40

This looks pretty dangerous.
Why not just use memref.transpose which has better static guarantees?

[nicolasvasilache]

Side note: I think we could also drop memref.view and just use this instead.
This would be a good starter task.

mlir/test/Dialect/MemRef/canonicalize.mlir

154

doesn't this fundamentally break the ability to do the "stride-0" trick?

[herhut]

Thanks!

mlir/test/Dialect/MemRef/canonicalize.mlir

154

We never do the stride-0 trick in the static case, so the 0 is always a dynamic value. Because we cannot express the static case properly. It would require an affine map with a 0 constant (otherwise the type is a lie).

The stride-0 trick if fishy anyway and we need to move away from it.

mlir/test/mlir-cpu-runner/copy.mlir

38–40

I prefer this form, as it makes it explicit what strides we have. This is just to test that the copy honors that information.

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

return op.emitError("different element types specified for source type ")
            << srcType << " and result memref type " << resultType;
 
   // Match sizes in result memref type and in static_sizes attribute.
   for (auto &en :
        llvm::enumerate(llvm::zip(resultType.getShape(),
                                  extractFromI64ArrayAttr(op.static_sizes())))) {
     int64_t resultSize = std::get<0>(en.value());
     int64_t expectedSize = std::get<1>(en.value());
-    if (!ShapedType::isDynamic(resultSize) && resultSize != expectedSize)
+    if (!ShapedType::isDynamic(resultSize) &&
+        !ShapedType::isDynamic(expectedSize) && resultSize != expectedSize)
       return op.emitError("expected result type with size = ")
              << expectedSize << " instead of " << resultSize
              << " in dim = " << en.index();
   }
 
-  // Match offset and strides in static_offset and static_strides attributes if
-  // result memref type has an affine map specified.
-  if (!resultType.getLayout().isIdentity()) {
-    int64_t resultOffset;
-    SmallVector<int64_t, 4> resultStrides;
-    if (failed(getStridesAndOffset(resultType, resultStrides, resultOffset)))
-      return failure();
+  // Match offset and strides in static_offset and static_strides attributes. If
+  // result memref type has no affine map specified, this will assume an
+  // identity layout.
+    return op.emitError(
+               "expected result type to have strided layout but found ")
+           << resultType;
 
-    // Match offset in result memref type and in static_offsets attribute.
-    int64_t expectedOffset =
-        extractFromI64ArrayAttr(op.static_offsets()).front();
-    if (!ShapedType::isDynamicStrideOrOffset(resultOffset) &&
-        resultOffset != expectedOffset)
-      return op.emitError("expected result type with offset = ")
-             << resultOffset << " instead of " << expectedOffset;
+  int64_t expectedOffset = extractFromI64ArrayAttr(op.static_offsets()).front();
+      !ShapedType::isDynamicStrideOrOffset(expectedOffset) &&
 
-    // Match strides in result memref type and in static_strides attribute.
-    for (auto &en : llvm::enumerate(llvm::zip(
-             resultStrides, extractFromI64ArrayAttr(op.static_strides())))) {
-      int64_t resultStride = std::get<0>(en.value());
-      int64_t expectedStride = std::get<1>(en.value());
-      if (!ShapedType::isDynamicStrideOrOffset(resultStride) &&
-          resultStride != expectedStride)
-        return op.emitError("expected result type with stride = ")
-               << expectedStride << " instead of " << resultStride
-               << " in dim = " << en.index();
-    }
-  }
+        !ShapedType::isDynamicStrideOrOffset(expectedStride) &&
+
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Reassociative reshape ops
 //===----------------------------------------------------------------------===//
 
 SmallVector<AffineMap, 4> CollapseShapeOp::getReassociationMaps() {

mlir/test/Dialect/MemRef/canonicalize.mlir

 // -----
 
 // CHECK-LABEL: func @dim_of_sized_view
 //  CHECK-SAME:   %{{[a-z0-9A-Z_]+}}: memref<?xi8>
 //  CHECK-SAME:   %[[SIZE:.[a-z0-9A-Z_]+]]: index
 //       CHECK:   return %[[SIZE]] : index
 func @dim_of_sized_view(%arg : memref<?xi8>, %size: index) -> index {
   %c0 = arith.constant 0 : index
-  %0 = memref.reinterpret_cast %arg to offset: [0], sizes: [%size], strides: [0] : memref<?xi8> to memref<?xi8>
+  %0 = memref.reinterpret_cast %arg to offset: [0], sizes: [%size], strides: [1] : memref<?xi8> to memref<?xi8>
   %1 = memref.dim %0, %c0 : memref<?xi8>
   return %1 : index
 }
 
 // -----
 
 // CHECK-LABEL: func @no_fold_of_store
 //  CHECK:   %[[cst:.+]] = memref.cast %arg

mlir/test/Dialect/MemRef/invalid.mlir

func @memref_reinterpret_cast_offset_mismatch(%in: memref<?xf32>) {
   %out = memref.reinterpret_cast %in to
            offset: [2], sizes: [10], strides: [2]
          : memref<?xf32> to memref<10xf32, offset: 2, strides: [1]>
   return
 }
 
 // -----
 
+func @memref_reinterpret_cast_no_map_but_offset(%in: memref<?xf32>) {
+  // expected-error @+1 {{expected result type with offset = 0 instead of 2}}
+  %out = memref.reinterpret_cast %in to offset: [2], sizes: [10], strides: [1]
+         : memref<?xf32> to memref<10xf32>
+  return
+}
+
+// -----
+
+func @memref_reinterpret_cast_no_map_but_stride(%in: memref<?xf32>) {
+  // expected-error @+1 {{expected result type with stride = 10 instead of 1 in dim = 0}}
+  %out = memref.reinterpret_cast %in to offset: [0], sizes: [10], strides: [10]
+         : memref<?xf32> to memref<10xf32>
+  return
+}
+
+// -----
+
+func @memref_reinterpret_cast_no_map_but_strides(%in: memref<?x?xf32>) {
+  // expected-error @+1 {{expected result type with stride = 42 instead of 10 in dim = 0}}
+  %out = memref.reinterpret_cast %in to
+           offset: [0], sizes: [9, 10], strides: [42, 1]
+         : memref<?x?xf32> to memref<9x10xf32>
+  return
+}
+
+// -----
+
+func @memref_reinterpret_cast_non_strided_layout(%in: memref<?x?xf32>) {
+  // expected-error @+1 {{expected result type to have strided layout but found 'memref<9x10xf32, affine_map<(d0, d1) -> (d0)>>}}
+  %out = memref.reinterpret_cast %in to
+           offset: [0], sizes: [9, 10], strides: [42, 1]
+         : memref<?x?xf32> to memref<9x10xf32, affine_map<(d0, d1) -> (d0)>>
+  return
+}
+
+// -----
+
 func @memref_reshape_element_type_mismatch(
        %buf: memref<*xf32>, %shape: memref<1xi32>) {
   // expected-error @+1 {{element types of source and destination memref types should be the same}}
   memref.reshape %buf(%shape) : (memref<*xf32>, memref<1xi32>) -> memref<?xi32>
 }
 
 // -----
 

mlir/test/Dialect/MemRef/ops.mlir

 func @memref_reinterpret_cast_static_to_dynamic_sizes(%in: memref<?xf32>)
     -> memref<10x?xf32, offset: ?, strides: [?, 1]> {
   %out = memref.reinterpret_cast %in to
            offset: [1], sizes: [10, 10], strides: [1, 1]
            : memref<?xf32> to memref<10x?xf32, offset: ?, strides: [?, 1]>
   return %out : memref<10x?xf32, offset: ?, strides: [?, 1]>
 }
 
+// CHECK-LABEL: func @memref_reinterpret_cast_dynamic_offset
+func @memref_reinterpret_cast_dynamic_offset(%in: memref<?xf32>, %offset: index)
+    -> memref<10x?xf32, offset: ?, strides: [?, 1]> {
+  %out = memref.reinterpret_cast %in to
+           offset: [%offset], sizes: [10, 10], strides: [1, 1]
+           : memref<?xf32> to memref<10x?xf32, offset: ?, strides: [?, 1]>
+  return %out : memref<10x?xf32, offset: ?, strides: [?, 1]>
+}
+
 // CHECK-LABEL: func @memref_reshape(
 func @memref_reshape(%unranked: memref<*xf32>, %shape1: memref<1xi32>,
          %shape2: memref<2xi32>, %shape3: memref<?xi32>) -> memref<*xf32> {
   %dyn_vec = memref.reshape %unranked(%shape1)
                : (memref<*xf32>, memref<1xi32>) -> memref<?xf32>
   %dyn_mat = memref.reshape %dyn_vec(%shape2)
                : (memref<?xf32>, memref<2xi32>) -> memref<?x?xf32>
   %new_unranked = memref.reshape %dyn_mat(%shape3)

mlir/test/mlir-cpu-runner/copy.mlir

func @main() -> () {
   memref.copy %input, %copy : memref<2x3xf32> to memref<2x3xf32>
   %unranked_copy = memref.cast %copy : memref<2x3xf32> to memref<*xf32>
   call @print_memref_f32(%unranked_copy) : (memref<*xf32>) -> ()
   // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1]
   // CHECK-NEXT: [0,   1,   2]
   // CHECK-NEXT: [3,   4,   5]
 
   %copy_two = memref.alloc() : memref<3x2xf32>
-  %copy_two_casted = memref.reinterpret_cast %copy_two to offset: [0], sizes: [2,3], strides:[1, 2]
-    : memref<3x2xf32> to memref<2x3xf32>
-  memref.copy %input, %copy_two_casted : memref<2x3xf32> to memref<2x3xf32>
+  %copy_two_casted = memref.reinterpret_cast %copy_two to offset: [0], sizes: [2, 3], strides:[1, 2]
+    : memref<3x2xf32> to memref<2x3xf32, offset: 0, strides: [1, 2]>
+  memref.copy %input, %copy_two_casted : memref<2x3xf32> to memref<2x3xf32, offset: 0, strides: [1, 2]>
   %unranked_copy_two = memref.cast %copy_two : memref<3x2xf32> to memref<*xf32>
   call @print_memref_f32(%unranked_copy_two) : (memref<*xf32>) -> ()
   // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1]
   // CHECK-NEXT: [0,   3]
   // CHECK-NEXT: [1,   4]
   // CHECK-NEXT: [2,   5]
 
   %input_empty = memref.alloc() : memref<3x0x1xf32>