[mlir][memref] ReinterpretCast: allow static sizes/strides/offset where affine map expects dynamic

[Hardcode84]

Summary

• There is no reason to forbid that case
• Also, user will get very unfriendly error like expected result type with offset = -9223372036854775808 instead of 1

[herhut]

Thanks. I don't see any reason to disallow this either and was surprised by the error message.

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 (resultSize != expectedSize)
+    if (!ShapedType::isDynamic(resultSize) && 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 in result memref type and in static_offsets attribute.
     int64_t expectedOffset =
         extractFromI64ArrayAttr(op.static_offsets()).front();
-    if (resultOffset != expectedOffset)
+    if (!ShapedType::isDynamicStrideOrOffset(resultOffset) &&
+        resultOffset != expectedOffset)
       return op.emitError("expected result type with offset = ")
              << resultOffset << " instead of " << 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 (resultStride != expectedStride)
+      if (!ShapedType::isDynamicStrideOrOffset(resultStride) &&
+          resultStride != expectedStride)
         return op.emitError("expected result type with stride = ")
                << expectedStride << " instead of " << resultStride
                << " in dim = " << en.index();
     }
   }
   return success();
 }
 

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_offset_mismatch(%in: memref<?xf32>) {
-  %c0 = arith.constant 0 : index
-  %c10 = arith.constant 10 : index
-  // expected-error @+1 {{expected result type with size = 10 instead of -1 in dim = 0}}
-  %out = memref.reinterpret_cast %in to
-           offset: [%c0], sizes: [10, %c10], strides: [%c10, 1]
-           : memref<?xf32> to memref<?x?xf32, offset: ?, strides: [?, 1]>
-  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(%in: memref<?xf32>)
   %c0 = arith.constant 0 : index
   %c10 = arith.constant 10 : index
   %out = memref.reinterpret_cast %in to
            offset: [%c0], sizes: [10, %c10], strides: [%c10, 1]
            : memref<?xf32> to memref<10x?xf32, offset: ?, strides: [?, 1]>
   return %out : memref<10x?xf32, offset: ?, strides: [?, 1]>
 }
 
+// CHECK-LABEL: func @memref_reinterpret_cast_static_to_dynamic_sizes
+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_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)