[mlir] Fold memref.reinterpret_cast(x) -> x when type does not change

[aturetsk]

Summary

Fold memref.reinterpret_cast(x) -> x when type does not change.

[aturetsk]

Ping

[aturetsk]

Ping

[aturetsk]

Ping

[aturetsk]

@nicolasvasilache, @mehdi_amini, @ThomasRaoux
Could anyone please take a look at this?

[aturetsk]

Ping

[nicolasvasilache]

My apologies, I completely missed these notifications, LGTM

[nicolasvasilache]

can you update the commit message to say: "Fold memref.reinterpret_cast(x) -> x when the type is fully static and does not change." ?

[aturetsk]

Thanks for review! Updated the commit message like you suggested.

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

auto getPrevSrc = [&]() -> Value {
     return nullptr;
   };
 
   if (auto prevSrc = getPrevSrc()) {
     getSourceMutable().assign(prevSrc);
     return getResult();
   }
 
+  // reinterpret_cast(x) w/o offset/shape/stride changes -> x
+  if (!ShapedType::isDynamicShape(getType().getShape()) &&
+      src.getType() == getType() && getStaticOffsets().front() == 0) {
+    return src;
+  }
+
   return nullptr;
 }
 
 SmallVector<OpFoldResult> ReinterpretCastOp::getConstifiedMixedSizes() {
   SmallVector<OpFoldResult> values = getMixedSizes();
   constifyIndexValues(values, getType(), getContext(), getConstantSizes,
                       ShapedType::isDynamic);
   return values;

mlir/test/Dialect/MemRef/canonicalize.mlir

func.func @scopeInline(%arg : memref<index>) {
   return
 }
 
 // CHECK:   func @scopeInline
 // CHECK-NOT:  memref.alloca_scope
 
 // -----
 
+// CHECK-LABEL: func @reinterpret_noop
+//  CHECK-SAME: (%[[ARG:.*]]: memref<2x3x4xf32>)
+//  CHECK-NEXT: return %[[ARG]]
+func.func @reinterpret_noop(%arg : memref<2x3x4xf32>) -> memref<2x3x4xf32> {
+  %0 = memref.reinterpret_cast %arg to offset: [0], sizes: [2, 3, 4], strides: [12, 4, 1] : memref<2x3x4xf32> to memref<2x3x4xf32>
+  return %0 : memref<2x3x4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @reinterpret_of_reinterpret
 //  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.func @reinterpret_of_reinterpret(%arg : memref<?xi8>, %size1: index, %size2: index) -> memref<?xi8> {
   %0 = memref.reinterpret_cast %arg to offset: [0], sizes: [%size1], strides: [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>