[mlir] Normalize memrefs in LoadOp and StoreOp of Standard Ops

[imaihal]

Summary

Added a trait, MemRefsNormalizable in LoadOp and StoreOp of Standard Ops
to normalize input memrefs in LoadOp and StoreOp.

Related revision: https://reviews.llvm.org/D86236

[imaihal]

When we use load and store ops like the test in this patch, the --normalize-memrefs does not work. This patch fix it by adding MemRefsNormalizable trait in load and store ops.
If we can use affine.load and affine.store in this test, the normalization works, but sometimes affine.load can not be used. In the case, we need to use std.load.

More practical example is here. https://github.com/llvm/llvm-project/blob/master/mlir/test/Dialect/Linalg/affine.mlir#L92-L94 affine.load can not be used here.

@bondhugula @AlexEichenberger @avarmapml Could you review this patch?

[bondhugula]

Thanks for adding this!

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

64

Nit: for readability, it may be good to define the map to use here right before the function instead of reusing #map0 defined all the way up. You could use say:

#map_tile = ...

[AlexEichenberger]

Thanks, glad to see that the traits make such changes easy... Please correct Uday's nit, and its good to go.

[imaihal]

Reflected the comment about map, and rebased

[imaihal]

@AlexEichenberger

glad to see that the traits make such changes easy...

Yes. The traits are very useful. We might need to add them in other ops depending on use-case, but it is easy.

@bondhugula I don't have commit permission. Could you commit this patch?

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 
 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 
 def LoadOp : Std_Op<"load",
      [TypesMatchWith<"result type matches element type of 'memref'",
                      "memref", "result",
-                     "$_self.cast<MemRefType>().getElementType()">]> {
+                     "$_self.cast<MemRefType>().getElementType()">,
+                     MemRefsNormalizable]> {
   let summary = "load operation";
   let description = [{
     The `load` op reads an element from a memref specified by an index list. The
     output of load is a new value with the same type as the elements of the
     memref. The arity of indices is the rank of the memref (i.e., if the memref
     loaded from is of rank 3, then 3 indices are required for the load following
     the memref identifier).
 

 
 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//
 
 def StoreOp : Std_Op<"store",
      [TypesMatchWith<"type of 'value' matches element type of 'memref'",
                      "memref", "value",
-                     "$_self.cast<MemRefType>().getElementType()">]> {
+                     "$_self.cast<MemRefType>().getElementType()">,
+                     MemRefsNormalizable]> {
   let summary = "store operation";
   let description = [{
     Store a value to a memref location given by indices. The value stored should
     have the same type as the elemental type of the memref. The number of
     arguments provided within brackets need to match the rank of the memref.
 
     In an affine context, the indices of a store are restricted to SSA values
     bound to surrounding loop induction variables,

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

func @test_norm_mix(%arg0 : memref<1x16x1x1x32x64xf32>) -> () {
     "test.op_norm"(%arg0, %0) : (memref<1x16x1x1x32x64xf32>, memref<1x16x14x14xf32, #map0>) -> ()
     dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
     // CHECK: %[[v0:[a-z0-9]*]] = alloc() : memref<1x16x1x1x32x64xf32>
     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
     // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
     return
 }
+
+// Test with maps in load and store ops.
+
+#map_tile = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 32, d2 mod 32, d3 mod 32)>
+
+// CHECK-LABEL: test_load_store
+// CHECK-SAME: (%[[ARG0:[a-z0-9]*]]: memref<1x16x14x14xf32>
+func @test_load_store(%arg0 : memref<1x16x14x14xf32>) -> () {
+    %0 = alloc() : memref<1x16x14x14xf32, #map_tile>
+    // CHECK: %[[v0:[a-z0-9]*]] = alloc() : memref<1x16x1x1x32x32xf32>
+    %1 = alloc() : memref<1x16x14x14xf32>
+    // CHECK: %[[v1:[a-z0-9]*]] = alloc() : memref<1x16x14x14xf32>
+    "test.op_norm"(%0, %1) : (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>) -> ()
+    // CHECK: "test.op_norm"(%[[v0]], %[[v1]]) : (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>) -> ()
+    %cst = constant 3.0 : f32
+    affine.for %i = 0 to 1 {
+      affine.for %j = 0 to 16 {
+        affine.for %k = 0 to 14 {
+          affine.for %l = 0 to 14 {
+            %2 = load %1[%i, %j, %k, %l] : memref<1x16x14x14xf32>
+            // CHECK: memref<1x16x14x14xf32>
+            %3 = addf %2, %cst : f32
+            store %3, %arg0[%i, %j, %k, %l] : memref<1x16x14x14xf32>
+            // CHECK: memref<1x16x14x14xf32>
+          }
+        }
+      }
+    }
+    dealloc %0 :  memref<1x16x14x14xf32, #map_tile>
+    // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
+    dealloc %1 :  memref<1x16x14x14xf32>
+    // CHECK: dealloc %[[v1]] : memref<1x16x14x14xf32>
+    return
+}