[mlir][memref] memref.reinterpret_cast folding

[Hardcode84]

Summary

• reinterpret_cast(reinterpret_cast(x)) -> reinterpret_cast(x)
• reinterpret_cast(cast(x)) -> reinterpret_cast(x)
• reinterpret_cast(subview(x)) -> reinterpret_cast(x) if subview offsets are 0

[Hardcode84]

fix style

[Hardcode84]

fix gcc build

[bondhugula]

These are all pretty useful - thanks!

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

1353–1378

Please add comments along the lines you have in the commit summary.

1363

Avoid suffixed names like is0.

[Hardcode84]

add comments

[Hardcode84]

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

1363

Sorry, I didn't quite get this comment. Renamed to isZeroVal anyway.

[bondhugula]

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

1363

isZero was what I meant instead of is0. This is fine now.

[Hardcode84]

ping

[nicolasvasilache]

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

1367

Could you add a isConstantIntValue(OpFoldResult, int64_t) to the utils and use it here?

1371

This would become:

if (llvm::all_of(prev.getMixedOffsets(), [](OpFoldResult val){ return isConstantIntValue(val, 0); }))
  ...

[Hardcode84]

rebase, review comments

[nicolasvasilache]

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

56 ↗

(On Diff #414649)

nit: typo

[nicolasvasilache]

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

84 ↗

(On Diff #414649)

nit: typo

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

std::array<unsigned, 3> getArrayAttrMaxRanks() {
       unsigned resultRank = getResult().getType().cast<ShapedType>().getRank();
       return {1, resultRank, resultRank};
     }
 
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
   }];
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_RankOp : MemRef_Op<"rank", [NoSideEffect]> {
   let summary = "rank operation";

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

if (!ShapedType::isDynamicStrideOrOffset(resultStride) &&
       return emitError("expected result type with stride = ")
              << expectedStride << " instead of " << resultStride
              << " in dim = " << en.index();
   }
 
   return success();
 }
 
+OpFoldResult ReinterpretCastOp::fold(ArrayRef<Attribute> operands) {
+  Value src = source();
+  auto getPrevSrc = [&]() -> Value {
+    if (auto prev = src.getDefiningOp<ReinterpretCastOp>())
+      return prev.source();
+
+    if (auto prev = src.getDefiningOp<CastOp>())
+      return prev.source();
+
+    if (auto prev = src.getDefiningOp<SubViewOp>()) {
+      auto is0 = [](OpFoldResult val) {
+        auto i = getConstantIntValue(val);
+        return i && *i == 0;
+      };
+      if (llvm::all_of(prev.getMixedOffsets(), is0))
+        return prev.source();
+    }
+
+    return nullptr;
+  };
+
+  if (auto prevSrc = getPrevSrc()) {
+    sourceMutable().assign(prevSrc);
+    return *this;
+  }
+
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // Reassociative reshape ops
 //===----------------------------------------------------------------------===//
 
 SmallVector<AffineMap, 4> CollapseShapeOp::getReassociationMaps() {
   return getSymbolLessAffineMaps(getReassociationExprs());
 }
 SmallVector<ReassociationExprs, 4> CollapseShapeOp::getReassociationExprs() {

mlir/test/Dialect/MemRef/canonicalize.mlir

 
 func @self_copy(%m1: memref<?xf32>) {
   memref.copy %m1, %m1 : memref<?xf32> to memref<?xf32>
   return
 }
 
 // CHECK-LABEL: func @self_copy
 //  CHECK-NEXT:   return
+
+// -----
+
+// 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 @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>
+}
+
+// -----
+
+// CHECK-LABEL: func @reinterpret_of_cast
+//  CHECK-SAME: (%[[ARG:.*]]: memref<?xi8>, %[[SIZE:.*]]: index)
+//       CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [%[[SIZE]]], strides: [1]
+//       CHECK: return %[[RES]]
+func @reinterpret_of_cast(%arg : memref<?xi8>, %size: index) -> memref<?xi8> {
+  %0 = memref.cast %arg : memref<?xi8> to memref<5xi8>
+  %1 = memref.reinterpret_cast %0 to offset: [0], sizes: [%size], strides: [1] : memref<5xi8> to memref<?xi8>
+  return %1 : memref<?xi8>
+}
+
+// -----
+
+// CHECK-LABEL: func @reinterpret_of_subview
+//  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>
+}