[mlir] Support normalizing memrefs with MemRef_ReinterpretCastOp

[imaihal]

Summary

This patch enables normalizing memrefs with MemRef_ReinterpretCastOp by
adding MemRefsNormalizable trait in the Op definition.

Signed-off-by: Haruki Imai <[e-mail]>

[imaihal]

@bondhugula I would like to support normalizing memrefs with MemRef_ReinterpretCastOp. As usual, this patch adds MemRefNormalizable trait. Could you review this patch?

[imaihal]

@bondhugula Sorry again. Could you be a reviewer of this patch?

[bondhugula]

mlir/test/Transforms/normalize-memrefs-ops.mlir

128–131

Is this test case really testing this change? The memref_reinterpret_cast has no uses of the memref with the tiled layout map.

[imaihal]

mlir/test/Transforms/normalize-memrefs-ops.mlir

128–131

Yes. Without this patch, memrefs with the tiled layout map are not normalized in this function. I think it is because we take conservative approach described here.

https://mlir.llvm.org/doxygen/NormalizeMemRefs_8cpp_source.html#l00143

Check whether all the uses of AllocOps, CallOps and function arguments of a function are either of dereferencing type or are uses in: DeallocOp, CallOp or ReturnOp. Only if these constraints are satisfied will the function become a candidate for normalization. We follow a conservative approach here wherein even if the non-normalizable memref is not a part of the function's argument or return type, we still label the entire function as non-normalizable. We assume external functions to be normalizable.

[bondhugula]

mlir/test/Transforms/normalize-memrefs-ops.mlir

128–131

I see.

[imaihal]

@bondhugula Thanks for your review! I don't have commit access. So, could you land this patch?

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

 
 //===----------------------------------------------------------------------===//
 // ReinterpretCastOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_ReinterpretCastOp:
     BaseOpWithOffsetSizesAndStrides<MemRef_Dialect, "reinterpret_cast", [
       NoSideEffect, AttrSizedOperandSegments, ViewLikeOpInterface,
-      OffsetSizeAndStrideOpInterface
+      OffsetSizeAndStrideOpInterface, MemRefsNormalizable
     ]> {
   let summary = "memref reinterpret cast operation";
   let description = [{
     Modify offset, sizes and strides of an unranked/ranked memref.
 
     Example:
     ```mlir
     memref.reinterpret_cast %ranked to

mlir/test/Transforms/normalize-memrefs-ops.mlir

func @test_norm_ret(%arg0: memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>) {
     // CHECK-NEXT: memref.dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
     return %1, %2 : memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>
     // CHECK-NEXT: return %[[v1]], %[[v2]] : memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>
 }
 
 // Test with an arbitrary op that references the function symbol.
 
 "test.op_funcref"() {func = @test_norm_mix} : () -> ()
+
+
+// -----
+
+#map_1d_tile = affine_map<(d0) -> (d0 floordiv 32, d0 mod 32)>
+
+// Test with memref.reinterpret_cast
+
+// CHECK-LABEL: test_norm_reinterpret_cast
+// CHECK-SAME: (%[[ARG0:.*]]: memref<1x32xf32>) -> memref<3x1x1xf32> {
+func @test_norm_reinterpret_cast(%arg0 : memref<3xf32, #map_1d_tile>) -> (memref<3x1x1xf32>) {
+    %0 = memref.alloc() : memref<3xf32>
+    "test.op_norm"(%arg0, %0) : (memref<3xf32, #map_1d_tile>, memref<3xf32>) -> ()
+    %1 = memref.reinterpret_cast %0 to offset: [0], sizes: [3, 1, 1], strides: [1, 1, 1] : memref<3xf32> to memref<3x1x1xf32>
+    // CHECK: %[[v0:.*]] = memref.alloc() : memref<3xf32>
+    // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x32xf32>, memref<3xf32>) -> ()
+    // CHECK: memref.reinterpret_cast %[[v0]] to offset: [0], sizes: [3, 1, 1], strides: [1, 1, 1] : memref<3xf32> to memref<3x1x1xf32>
+    return %1 : memref<3x1x1xf32>
+}