[mlir] Normalize the results of normalizable operations

[imaihal]

Summary

Memrefs with affine_map in the results of normalizable operation were
not normalized by --normalize-memrefs option. This patch normalizes
them.

[imaihal]

Fix to remove warning message (newOp -> *newOp)

[imaihal]

@bondhugula @avarmapml I found the results(returend memref) of normalizable operation were not normalized by --normalize-memrefs.
In test case, %1 was not normalized. This patch fixes it.
%1 = "test.op_norm_ret"(%arg0, %arg1) : (memref<1x32768xf32>, memref<1x16x14x14xf32, #map1>) -> (memref<1x16x14x14xf32, #map1>)

Could you review this patch?

[avarmapml]

Hi @imaihal, thank you for this patch.
The function normalizeFuncOpMemRefs has become quite long and I believe it'd be better to split the logic into modular sub-functions.
AllocOps, function arguments, result types, and external functions are being addressed within and maybe we can have separate functions for each of them.
Not sure if this change should go into this patch, maybe @bondhugula can say better.

mlir/lib/Transforms/NormalizeMemRefs.cpp

396

Comment should be terminated with a full stop.
Also, I guess the comment needs to be rephrased to Add types to resultTypes``?

[bondhugula]

mlir/lib/Transforms/NormalizeMemRefs.cpp

390–454

This is a really long lambda. Please outline.

[bondhugula]

Thanks for adding this support. Some minor comments. I'll be able to review post the refactoring.

mlir/lib/Transforms/NormalizeMemRefs.cpp

391

You don't need dyn_cast here - !isa<CallOp>(op) will do.

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

95

Nit: consider using map_tile.

mlir/test/lib/Dialect/Test/TestOps.td

630

-> ... of an op that has normalizable memref results.

[bondhugula]

mlir/lib/Transforms/NormalizeMemRefs.cpp

391

Nit: prefer using > 0.

432

If you do a dyn_cast, you have to check for null. If you know these will be memrefs, please use cast.

[imaihal]

Reflected comments and rebased.

[imaihal]

Updated test case to include two operation results.

[imaihal]

Rebased and fixed a comment.

[imaihal]

@bondhugula @avarmapml Thanks for your comments! I updated the code to reflect your comments. I moved part of code in the lambda to Util.cpp. Could you review again?

mlir/lib/Transforms/NormalizeMemRefs.cpp

390–454

I moved part of the code to newly created function in Util.cpp. So, I think the lambda is not so long now. Also, I wrote comments about the outline at the top.

396

I wrote comments more in newly created function in Util.cpp.

[bondhugula]

mlir/include/mlir/Transforms/Utils.h

149–150 ↗

(On Diff #296336)

This has just one user and functionality specific to normalize-memrefs. Why is it exposed?

[imaihal]

Reflected the comment.

[imaihal]

mlir/include/mlir/Transforms/Utils.h

149–150 ↗

(On Diff #296336)

I moved the function to NormalizedMemrefs

[bondhugula]

Looks good. Please address the remaining minor comments.

mlir/lib/Transforms/NormalizeMemRefs.cpp

494–498

Triple doc comments please.

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

99

%[[v0:.*]] = should be sufficient.

[imaihal]

Rebased and reflected comments.

[imaihal]

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

99

I also fixed other variables from [a-z0-9]* to .* in this file.

[imaihal]

In D88719#2381285, @bondhugula wrote:

Looks good. Please address the remaining minor comments.

@bondhugula Thanks for your review. I reflected your comments. Could you commit this patch if you don't have any other comments? I don't have commit access.

mlir/lib/Transforms/NormalizeMemRefs.cpp

if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newMemRef,
       continue;
     }
 
     // All uses for the argument with old memref type were replaced
     // successfully. So we remove the old argument now.
     funcOp.front().eraseArgument(argIndex + 1);
   }
 
+  // Walk over normalizable operations to normalize memrefs of the operation
+  // results.
+  funcOp.walk([&](Operation *op) {
+    if (op->hasTrait<OpTrait::MemRefsNormalizable>() &&
+        op->getNumResults() >= 1 && !dyn_cast<CallOp>(op)) {
+      // Insert newOp containing normalized memref in the operation result.
+      OperationState result(op->getLoc(), op->getName());
+      result.addOperands(op->getOperands());
+      result.addAttributes(op->getAttrs());
+      // Add types to `result`
+      SmallVector<Type, 4> resultTypes;
+      bool resultTypeNormalized = false;
+      for (unsigned resIndex : llvm::seq<unsigned>(0, op->getNumResults())) {
+        auto resultType = op->getResult(resIndex).getType();
+        MemRefType memrefType = resultType.dyn_cast<MemRefType>();
+        // Check whether the operation result is MemRef type.
+        if (!memrefType) {
+          resultTypes.push_back(resultType);
+          continue;
+        }
+        // Fetch a new memref type after normalizing the old memref.
+        MemRefType newMemRefType =
+            normalizeMemRefType(memrefType, b,
+                                /*numSymbolicOperands=*/0);
+        if (newMemRefType == memrefType || funcOp.isExternal()) {
+          // Either memrefType already had an identity map or the map couldn't
+          // be transformed to an identity map.
+          resultTypes.push_back(memrefType);
+          continue;
+        }
+        resultTypes.push_back(newMemRefType);
+        resultTypeNormalized = true;
+      }
+      if (resultTypeNormalized) {
+        result.addTypes(resultTypes);
+        OpBuilder bb(op);
+        auto newOp = bb.createOperation(result);
+        bool replacingMemRefUsesFailed = false;
+
+        for (unsigned resIndex : llvm::seq<unsigned>(0, op->getNumResults())) {
+          // Replace all uses of the old memrefs.
+          Value oldMemRef = op->getResult(resIndex);
+          Value newMemRef = newOp->getResult(resIndex);
+          MemRefType oldMemRefType = oldMemRef.getType().dyn_cast<MemRefType>();
+          MemRefType newMemRefType = newMemRef.getType().dyn_cast<MemRefType>();
+          if (oldMemRefType == newMemRefType)
+            continue;
+          // TODO: Multiple maps not supported
+          AffineMap layoutMap = oldMemRefType.getAffineMaps().front();
+          if (failed(replaceAllMemRefUsesWith(oldMemRef,
+                                              /*newMemRef=*/newMemRef,
+                                              /*extraIndices=*/{},
+                                              /*indexRemap=*/layoutMap,
+                                              /*extraOperands=*/{},
+                                              /*symbolOperands=*/{},
+                                              /*domInstFilter=*/nullptr,
+                                              /*postDomInstFilter=*/nullptr,
+                                              /*allowDereferencingOps=*/true,
+                                              /*replaceInDeallocOp=*/true))) {
+            newOp->erase();
+            replacingMemRefUsesFailed = true;
+            continue;
+          }
+        }
+        if (!replacingMemRefUsesFailed) {
+          // Replace old operation with new one and delete the old one.
+          op->replaceAllUsesWith(newOp);
+          op->erase();
+        }
+      }
+    }
+  });
+
   // In a normal function, memrefs in the return type signature gets normalized
   // as a result of normalization of functions arguments, AllocOps or CallOps'
   // result types. Since an external function doesn't have a body, memrefs in
   // the return type signature can only get normalized by iterating over the
   // individual return types.
   if (funcOp.isExternal()) {
     SmallVector<Type, 4> resultTypes;
     for (unsigned resIndex :
if (funcOp.isExternal()) {
 
     FunctionType newFuncType = FunctionType::get(/*inputs=*/inputTypes,
                                                  /*results=*/resultTypes,
                                                  /*context=*/&getContext());
     // Setting the new function signature for this external function.
     funcOp.setType(newFuncType);
   }
   updateFunctionSignature(funcOp, moduleOp);
 }

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

func @test_load_store(%arg0 : memref<1x16x14x14xf32>) -> () {
     }
     dealloc %0 :  memref<1x16x14x14xf32, #map_tile>
     // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
     dealloc %1 :  memref<1x16x14x14xf32>
     // CHECK: dealloc %[[v1]] : memref<1x16x14x14xf32>
     return
 }
 
+// Test with op_norm_ret, with maps in the results of normalizable operation.
+
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 32, d2 mod 32, d3 mod 32)>
+
+// CHECK-LABEL: test_norm_ret
+// CHECK-SAME: , %[[ARG1:[a-z0-9]*]]: memref<1x16x1x1x32x32xf32>) -> memref<1x16x1x1x32x32xf32> {
+func @test_norm_ret(%arg0: memref<1x32768xf32>, %arg1: memref<1x16x14x14xf32, #map1>) -> memref<1x16x14x14xf32, #map1> {
+    %0 = alloc() : memref<1x16x14x14xf32, #map1>
+    // CHECK-NEXT: %[[v0:[a-z0-9]*]] = alloc() : memref<1x16x1x1x32x32xf32>
+    %1 = "test.op_norm_ret"(%arg0, %arg1) : (memref<1x32768xf32>, memref<1x16x14x14xf32, #map1>) -> (memref<1x16x14x14xf32, #map1>)
+    // CHECK-NEXT: %[[v1:[a-z0-9]*]] = "test.op_norm_ret"
+    // CHECK-SAME: , memref<1x16x1x1x32x32xf32>) -> memref<1x16x1x1x32x32xf32>
+    "test.op_norm"(%1, %0) : (memref<1x16x14x14xf32, #map1>, memref<1x16x14x14xf32, #map1>) -> ()
+    // CHECK-NEXT: "test.op_norm"
+    // CHECK-SAME: : (memref<1x16x1x1x32x32xf32>, memref<1x16x1x1x32x32xf32>) -> ()
+    dealloc %0 : memref<1x16x14x14xf32, #map1>
+    // CHECK-NEXT: dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
+    return %1 : memref<1x16x14x14xf32, #map1>
+    // CHECK-NEXT: return %[[v1]] : memref<1x16x1x1x32x32xf32>
+}
+
 // Test with an arbitrary op that references the function symbol.
 
 "test.op_funcref"() {func = @test_norm_mix} : () -> ()

mlir/test/lib/Dialect/Test/TestOps.td

 // Test for memrefs normalization of an op with normalizable memrefs.
 def OpNorm : TEST_Op<"op_norm", [MemRefsNormalizable]> {
   let arguments = (ins AnyMemRef:$X, AnyMemRef:$Y);
 }
 // Test for memrefs normalization of an op without normalizable memrefs.
 def OpNonNorm : TEST_Op<"op_nonnorm"> {
   let arguments = (ins AnyMemRef:$X, AnyMemRef:$Y);
 }
+// Test for memrefs normalization of an op with normalizable memrefs.
+def OpNormRet : TEST_Op<"op_norm_ret", [MemRefsNormalizable]> {
+  let arguments = (ins AnyMemRef:$X, AnyMemRef:$Y);
+  let results = (outs AnyMemRef:$Z);
+}
 
 // Test for memrefs normalization of an op with a reference to a function
 // symbol.
 def OpFuncRef : TEST_Op<"op_funcref"> {
   let summary = "Test op with a reference to a function symbol";
   let description = [{
     The "test.op_funcref" is a test op with a reference to a function symbol.
   }];