[mlir][normalize-memrefs] Non-normalizable operations with identity map layouts do not block normalization of the entire function

[tungld]

Summary

The current approach is convervative in which whenever there is a non-normalizable operations in a function will the function be labelled as non-normalizable. It means it requires that all operations must have MemRefsNormalizable trait.

This patch relaxes the requirement that if the memref map layouts of a non-normalizable operation are identity, this operation does not block the normalization of the other operations in the same function.

[bondhugula]

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

163–164

This isn't the right to check for an identity layout. Instead, use isIdentity() on the layout itself.

178–179

Likewise: here and below.

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

49

Is this the right test case? In this case, all users of %arg0 are already normalizable and %arg0 doesn't have an identity layout map. Am I missing something?

[tungld]

• Use isIdentity() on the layout itself

[tungld]

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

163–164

Done. Thanks!

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

49

Yes, it is. Without this patch, %arg0 will not be normalized and the whole function stays unchanged. That's because the pass detects that the op_nonnorm op does not have MemRefsNormalizable trait no matter the op's arguments' map layouts are identity.

[tungld]

@bondhugula Is the current change ok? or should I make any other change? Thanks!

[tungld]

@bondhugula @avarmapml @imaihal could you please review this patch again? Thanks!

[bondhugula]

LGTM - thanks.

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

160–161

You can use allocOp.getType() to avoid using cast.

176–178

Use dyn_cast instead to avoid a repeated check.

191–192

Use dyn_cast.

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

for (func::FuncOp &funcOp : normalizableFuncs) {
       }
     }
   });
 }
 
 /// 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.
+/// become a candidate for normalization. When the uses of a memref are
+/// non-normalizable and the memref map layout is trivial (identity), we can
+/// still label the entire function as normalizable. We assume external
+/// functions to be normalizable.
 bool NormalizeMemRefs::areMemRefsNormalizable(func::FuncOp funcOp) {
   // We assume external functions to be normalizable.
   if (funcOp.isExternal())
     return true;
 
   if (funcOp
           .walk([&](memref::AllocOp allocOp) -> WalkResult {
             Value oldMemRef = allocOp.getResult();
-            if (!isMemRefNormalizable(oldMemRef.getUsers()))
+            if (!oldMemRef.getType()
+                     .cast<MemRefType>()
+                     .getLayout()
+                     .isIdentity() &&
+                !isMemRefNormalizable(oldMemRef.getUsers()))
               return WalkResult::interrupt();
             return WalkResult::advance();
           })
           .wasInterrupted())
     return false;
 
   if (funcOp
           .walk([&](func::CallOp callOp) -> WalkResult {
             for (unsigned resIndex :
                  llvm::seq<unsigned>(0, callOp.getNumResults())) {
               Value oldMemRef = callOp.getResult(resIndex);
               if (oldMemRef.getType().isa<MemRefType>())
-                if (!isMemRefNormalizable(oldMemRef.getUsers()))
+                if (!oldMemRef.getType()
+                         .cast<MemRefType>()
+                         .getLayout()
+                         .isIdentity() &&
+                    !isMemRefNormalizable(oldMemRef.getUsers()))
                   return WalkResult::interrupt();
             }
             return WalkResult::advance();
           })
           .wasInterrupted())
     return false;
 
   for (unsigned argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
     BlockArgument oldMemRef = funcOp.getArgument(argIndex);
     if (oldMemRef.getType().isa<MemRefType>())
-      if (!isMemRefNormalizable(oldMemRef.getUsers()))
+      if (!oldMemRef.getType().cast<MemRefType>().getLayout().isIdentity() &&
+          !isMemRefNormalizable(oldMemRef.getUsers()))
         return false;
   }
 
   return true;
 }
 
 /// Fetch the updated argument list and result of the function and update the
 /// function signature. This updates the function's return type at the caller

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

func.func @test_nonnorm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
     memref.dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
     // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x14x14xf32, #map>
     // CHECK: "test.op_nonnorm"(%[[ARG0]], %[[v0]]) : (memref<1x16x14x14xf32, #map>, memref<1x16x14x14xf32, #map>) -> ()
     // CHECK: memref.dealloc %[[v0]] : memref<1x16x14x14xf32, #map>
     return
 }
 
+// Test with op_nonnorm whose memref map layouts are identity. This op_nonnorm
+// does not block the normalization of other operations.
+
+// CHECK-LABEL: test_nonnorm_identity_layout
+// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>)
+func.func @test_nonnorm_identity_layout(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
+    %0 = memref.alloc() : memref<1x16x14x14xf32>
+    "test.op_nonnorm"(%0, %0) : (memref<1x16x14x14xf32>, memref<1x16x14x14xf32>) -> ()
+    "test.op_norm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32>) -> ()
+    memref.dealloc %0 :  memref<1x16x14x14xf32>
+
+    // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x14x14xf32>
+    // CHECK: "test.op_nonnorm"(%[[v0]], %[[v0]]) : (memref<1x16x14x14xf32>, memref<1x16x14x14xf32>) -> ()
+    // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x14x14xf32>) -> ()
+    // CHECK: memref.dealloc %[[v0]] : memref<1x16x14x14xf32>
+    return
+}
+
 // Test with op_norm, with maps in the operations in the function.
 
 // CHECK-LABEL: test_norm_mix
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>
 func.func @test_norm_mix(%arg0 : memref<1x16x1x1x32x64xf32>) -> () {
     %0 = memref.alloc() : memref<1x16x14x14xf32, #map0>
     "test.op_norm"(%arg0, %0) : (memref<1x16x1x1x32x64xf32>, memref<1x16x14x14xf32, #map0>) -> ()
     memref.dealloc %0 :  memref<1x16x14x14xf32, #map0>