[Mlir-commits] [mlir] b34b1c6 - [mlir] Support normalizing memrefs with MemRef_ReinterpretCastOp

[Uday Bondhugula]

Author: Haruki Imai
Date: 2021-08-11T01:15:18+05:30
New Revision: b34b1c695509747765de8bf3fde999fd6ec88367

URL: https://github.com/llvm/llvm-project/commit/b34b1c695509747765de8bf3fde999fd6ec88367
DIFF: https://github.com/llvm/llvm-project/commit/b34b1c695509747765de8bf3fde999fd6ec88367.diff

LOG: [mlir] Support normalizing memrefs with MemRef_ReinterpretCastOp

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

Signed-off-by: Haruki Imai <imaihal at jp.ibm.com>

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D107425

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
    mlir/test/Transforms/normalize-memrefs-ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 1ca01b8e9732c..c6c3ac8e12669 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -778,7 +778,7 @@ def MemRef_PrefetchOp : MemRef_Op<"prefetch"> {
 def MemRef_ReinterpretCastOp:
     BaseOpWithOffsetSizesAndStrides<MemRef_Dialect, "reinterpret_cast", [
       NoSideEffect, AttrSizedOperandSegments, ViewLikeOpInterface,
-      OffsetSizeAndStrideOpInterface
+      OffsetSizeAndStrideOpInterface, MemRefsNormalizable
     ]> {
   let summary = "memref reinterpret cast operation";
   let description = [{

diff  --git a/mlir/test/Transforms/normalize-memrefs-ops.mlir b/mlir/test/Transforms/normalize-memrefs-ops.mlir
index 01c694c62e1b6..afb2245e9b771 100644
--- a/mlir/test/Transforms/normalize-memrefs-ops.mlir
+++ b/mlir/test/Transforms/normalize-memrefs-ops.mlir
@@ -112,3 +112,22 @@ func @test_norm_ret(%arg0: memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14
 // 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>
+}