[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

397

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

391–455

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

392

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

669

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

[bondhugula]

mlir/lib/Transforms/NormalizeMemRefs.cpp

392

Nit: prefer using > 0.

433

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

391–455

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.

397

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

475–479

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

 /// non-normalizable as well. We assume external functions to be normalizable.
 struct NormalizeMemRefs : public NormalizeMemRefsBase<NormalizeMemRefs> {
   void runOnOperation() override;
   void normalizeFuncOpMemRefs(FuncOp funcOp, ModuleOp moduleOp);
   bool areMemRefsNormalizable(FuncOp funcOp);
   void updateFunctionSignature(FuncOp funcOp, ModuleOp moduleOp);
   void setCalleesAndCallersNonNormalizable(FuncOp funcOp, ModuleOp moduleOp,
                                            DenseSet<FuncOp> &normalizableFuncs);
+  Operation *createOpResultsNormalized(FuncOp funcOp, Operation *oldOp);
 };
 
 } // end anonymous namespace
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createNormalizeMemRefsPass() {
   return std::make_unique<NormalizeMemRefs>();
 }
 
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. When `op` has memrefs with affine map in the operation results,
+  // new operation containin normalized memrefs is created. Then, the memrefs
+  // are replaced. `CallOp` is skipped here because it is handled in
+  // `updateFunctionSignature()`.
+  funcOp.walk([&](Operation *op) {
+    if (op->hasTrait<OpTrait::MemRefsNormalizable>() &&
+        op->getNumResults() > 0 && !isa<CallOp>(op) && !funcOp.isExternal()) {
+      // Create newOp containing normalized memref in the operation result.
+      Operation *newOp = createOpResultsNormalized(funcOp, op);
+      // When all of the operation results have no memrefs or memrefs without
+      // affine map, `newOp` is the same with `op` and following process is
+      // skipped.
+      if (op != newOp) {
+        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>();
+          // Check whether the operation result is MemRef type.
+          if (!oldMemRefType)
+            continue;
+          MemRefType newMemRefType = newMemRef.getType().cast<MemRefType>();
+          if (oldMemRefType == newMemRefType)
+            continue;
+          // TODO: Assume single layout map. 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 other ops with new op and delete the old op when the
+          // replacement succeeded.
+          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 :
          llvm::seq<unsigned>(0, functionType.getNumResults())) {
       Type resType = functionType.getResult(resIndex);
       MemRefType memrefType = resType.dyn_cast<MemRefType>();
       // Check whether result is of MemRef type. Any other argument type can
       // simply be part of the final function signature.
       if (!memrefType) {
         resultTypes.push_back(resType);
         continue;
       }
       // Computing a new memref type after normalizing the old memref to have an
       // identity map layout.
       MemRefType newMemRefType = normalizeMemRefType(memrefType, b,
                                                      /*numSymbolicOperands=*/0);
       resultTypes.push_back(newMemRefType);
       continue;
     }
 
     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);
 }
+
+/// Create an operation containing normalized memrefs in the operation results.
+/// When the results of `oldOp` have memrefs with affine map, the memrefs are
+/// normalized, and new operation containing them in the operation results is
+/// returned. If all of the results of `oldOp` have no memrefs or memrefs
+/// without affine map, `oldOp` is returned without modification.
+Operation *NormalizeMemRefs::createOpResultsNormalized(FuncOp funcOp,
+                                                       Operation *oldOp) {
+  // Prepare OperationState to create newOp containing normalized memref in
+  // the operation results.
+  OperationState result(oldOp->getLoc(), oldOp->getName());
+  result.addOperands(oldOp->getOperands());
+  result.addAttributes(oldOp->getAttrs());
+  // Add normalized MemRefType to the OperationState.
+  SmallVector<Type, 4> resultTypes;
+  OpBuilder b(funcOp);
+  bool resultTypeNormalized = false;
+  for (unsigned resIndex : llvm::seq<unsigned>(0, oldOp->getNumResults())) {
+    auto resultType = oldOp->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) {
+      // 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;
+  }
+  result.addTypes(resultTypes);
+  // When all of the results of `oldOp` have no memrefs or memrefs without
+  // affine map, `oldOp` is returned without modification.
+  if (resultTypeNormalized) {
+    OpBuilder bb(oldOp);
+    return bb.createOperation(result);
+  } else
+    return oldOp;
+}

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

 // RUN: mlir-opt -normalize-memrefs %s | FileCheck %s
 
 // For all these cases, we test if MemRefs Normalization works with the test
 // operations.
 // * test.op_norm: this operation has the MemRefsNormalizable attribute. The tests
 //   that include this operation are constructed so that the normalization should
 //   happen.
 // * test_op_nonnorm: this operation does not have the MemRefsNormalization
 //   attribute. The tests that include this operation are contructed so that the
 //    normalization should not happen.
 
 #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 64, d2 mod 32, d3 mod 64)>
 
 // Test with op_norm and maps in arguments and in the operations in the function.
 
 // CHECK-LABEL: test_norm
-// CHECK-SAME: (%[[ARG0:[a-z0-9]*]]: memref<1x16x1x1x32x64xf32>)
+// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>)
 func @test_norm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
     %0 = alloc() : memref<1x16x14x14xf32, #map0>
     "test.op_norm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
     dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
-    // CHECK: %[[v0:[a-z0-9]*]] = alloc() : memref<1x16x1x1x32x64xf32>
+    // CHECK: %[[v0:.*]] = alloc() : memref<1x16x1x1x32x64xf32>
     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
     // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
     return
 }
 
 // Same test with op_nonnorm, with maps in the arguments and the operations in the function.
 
 // CHECK-LABEL: test_nonnorm
-// CHECK-SAME: (%[[ARG0:[a-z0-9]*]]: memref<1x16x14x14xf32, #map>)
+// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32, #map>)
 func @test_nonnorm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
     %0 = alloc() : memref<1x16x14x14xf32, #map0>
     "test.op_nonnorm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
     dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
-    // CHECK: %[[v0:[a-z0-9]*]] = alloc() : memref<1x16x14x14xf32, #map>
+    // CHECK: %[[v0:.*]] = alloc() : memref<1x16x14x14xf32, #map>
     // CHECK: "test.op_nonnorm"(%[[ARG0]], %[[v0]]) : (memref<1x16x14x14xf32, #map>, memref<1x16x14x14xf32, #map>) -> ()
     // CHECK: dealloc %[[v0]] : memref<1x16x14x14xf32, #map>
     return
 }
 
 // Test with op_norm, with maps in the operations in the function.
 
 // CHECK-LABEL: test_norm_mix
-// CHECK-SAME: (%[[ARG0:[a-z0-9]*]]: memref<1x16x1x1x32x64xf32>
+// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>
 func @test_norm_mix(%arg0 : memref<1x16x1x1x32x64xf32>) -> () {
     %0 = alloc() : memref<1x16x14x14xf32, #map0>
     "test.op_norm"(%arg0, %0) : (memref<1x16x1x1x32x64xf32>, memref<1x16x14x14xf32, #map0>) -> ()
     dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
-    // CHECK: %[[v0:[a-z0-9]*]] = alloc() : memref<1x16x1x1x32x64xf32>
+    // CHECK: %[[v0:.*]] = 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>
+// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32>
 func @test_load_store(%arg0 : memref<1x16x14x14xf32>) -> () {
     %0 = alloc() : memref<1x16x14x14xf32, #map_tile>
-    // CHECK: %[[v0:[a-z0-9]*]] = alloc() : memref<1x16x1x1x32x32xf32>
+    // CHECK: %[[v0:.*]] = alloc() : memref<1x16x1x1x32x32xf32>
     %1 = alloc() : memref<1x16x14x14xf32>
-    // CHECK: %[[v1:[a-z0-9]*]] = alloc() : memref<1x16x14x14xf32>
+    // CHECK: %[[v1:.*]] = 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
 }
 
+// Test with op_norm_ret, with maps in the results of normalizable operation.
+
+// CHECK-LABEL: test_norm_ret
+// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x32xf32>) -> (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>) {
+func @test_norm_ret(%arg0: memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>) {
+    %0 = alloc() : memref<1x16x14x14xf32, #map_tile>
+    // CHECK-NEXT: %[[v0:.*]] = alloc() : memref<1x16x1x1x32x32xf32>
+    %1, %2 = "test.op_norm_ret"(%arg0) : (memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>)
+    // CHECK-NEXT: %[[v1:.*]], %[[v2:.*]] = "test.op_norm_ret"
+    // CHECK-SAME: (memref<1x16x1x1x32x32xf32>) -> (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>)
+    "test.op_norm"(%1, %0) : (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32, #map_tile>) -> ()
+    // CHECK-NEXT: "test.op_norm"
+    // CHECK-SAME: : (memref<1x16x1x1x32x32xf32>, memref<1x16x1x1x32x32xf32>) -> ()
+    dealloc %0 : memref<1x16x14x14xf32, #map_tile>
+    // CHECK-NEXT: 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} : () -> ()

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 that has normalizable memref results.
+def OpNormRet : TEST_Op<"op_norm_ret", [MemRefsNormalizable]> {
+  let arguments = (ins AnyMemRef:$X);
+  let results = (outs AnyMemRef:$Y, 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.
   }];