[MLIR] Introduce inter-procedural memref layout normalization

[avarmapml]

Summary

• Introduces a pass that normalizes the affine layout maps to the identity layout map both within and across functions by rewriting function arguments and call operands where necessary.
• Memref normalization is now implemented entirely in the module pass '-normalize-memrefs' and the limited intra-procedural version has been removed from '-simplify-affine-structures'.
• Run using -normalize-memrefs.
• Return ops are not handled and would be handled in the subsequent revisions.

Signed-off-by: Abhishek Varma <[e-mail]>

[avarmapml]

Reformatted the code using clang-format

• fixed pre-merge checks' clang-format refactoring error.

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

115–116

You need to use a cast instead of static_cast here. Anyway, you could just use a dyn_cast to avoid doing isa first and then cast.

[bondhugula]

mlir/include/mlir/Transforms/Utils.h

90–92

It's not clear what 'trivializing" is. Also, the name 'getNewMemRefType' isn't descriptive enough / suitable. Also, this doesn't need to be perhaps exposed.

[avarmapml]

Code refactoring:

• Used dyn_cast for CallOp.
• Changed name of utility function which helps transform memref map layout.

[bondhugula]

mlir/include/mlir/Transforms/Passes.h

26

Nit: It'd be good to leave a blank line here in between.

mlir/include/mlir/Transforms/Utils.h

90

This is really not just transforming the memref layout map but transforming the memref type so that the shape and layout map are modified with the layout map being normalized to identity. This doc comment will have to rewritten. Also, this should be called something like normalizeMemRefType or normalizeMemRefLayout? 'transform' doesn't say how it's being transformed.

mlir/lib/Transforms/NormalizeMemRefs.cpp

10

-> ... to normalize memrefs to have identity layout maps?

33–35

Since these are doc comments, perhaps, it's better to move this entire comment block to right above the struct NormalizeMemRefs.

173

You don't need the static_cast.

180–181

You can just use argTypes and resultsTypes here - they will implicitly convert to ArrayRef.

189

Don't need the parentheses.

202

cast<CallOp>(callOp)

216–217

Looks like this will just crash if your call op is returning something that isn't a memref. So you are also missing a test case to exercise this?

217

dyn_cast -> cast if you already know it's a memref type. Otherwise, you'll have to check where the dyn_cast returned null!

mlir/lib/Transforms/Utils/Utils.cpp

100–124

You actually don't need all of this. Instead, just setOperand will work, i.e., setOperand(memRefOperandPos, newMemRef).

[avarmapml]

Review changes :-

• Added a test case to check functions which return values of non-memref types.
• Function renaming and doc-comment changes.
• Minor code changes.

[avarmapml]

Using setOperand to change memref operand of an operation.

[rriddle]

mlir/include/mlir/Transforms/Passes.td

409

Please keep these sorted alphabetically.

[rriddle]

mlir/lib/Transforms/NormalizeMemRefs.cpp

2

This line should be 80 lines long.

20

Can you trim these includes? Some shouldn't be necessary, e.g. Operation.h and FunctionSupport shouldn't be necessary.

70

nit: Please add variable names to declarations as well.

80

nit: I don't think this comment is necessary/useful.

86

nit: Please remove the user name from the TODO.

90

This is a very large function, can you split it up?

96

nit: Drop the trivial braces here and below.

104

nit: Can you use a longer name than ft?

110

nit: This could iterate over ft.getResults(), or use something like llvm::to_vector<4>(ft.getResults()).

116

Can you just use llvm::enumerate(funcOp.getArguments()) instead? It provides you with a value+index pair.

121

nit: Can you merge this comment with the one before the if?

140

insertArgument returns the inserted argument.

154

You aren't adding anything to argTypes, is that intended?

161

nit: Remove the user name here and other places.

163

nit: Can you iterate resultTypes directly instead?

for (Type &resultType : resultTypes)
  if (resultType == memrefType)
    resultType = newMemRef.getType();

181

nit: Remove the extra * before context.

185

funcOp.getSymbolUses(moduleOp)?

195

Please do not use OperationState directly, use builder.create<CallOp> instead.

231

This looks incorrect, if any of the argument types changed you can't keep the same call op.

mlir/lib/Transforms/Utils/Utils.cpp

90

Same here, please avoid using OperationState.

439

nit: Remove the wrapping ().

[bondhugula]

It would be best to handle return ops (where they return memref values with non-identity layouts to be normalized) separately in the next revision - since they add significant complexity.

[avarmapml]

Memref map layout normalization for function argument and AllocOps

• Support for ReturnOp will be handled in the next revision. So now we only deal with function arguments and allocOps whose uses are :-
  1. Dereferencing types
  2. Non-dereferencing types : DeallocOp, CallOp
• Code review changes.
• Added test cases.

[avarmapml]

mlir/lib/Transforms/NormalizeMemRefs.cpp

195

So, the ReturnOps will be handled in the next revision and I'll use this there.
Thank you!

231

I understand. The next revision will take care of this.

mlir/lib/Transforms/Utils/Utils.cpp

90

Unlike the previous use of OperationState where it was known to be CallOp here it can be of any dereferencing type, so OpTy for builder.create<OpTy> is unknown.

[avarmapml]

Reformatted test cases

[avarmapml]

Iterating over function arguments

• llvm::enumerate didn't work on MutableArrayRef in the previous build.

[bondhugula]

Could you update the commit summary on this revision page to update the return op handing part? (It doesn't pick it up from arc - you'll have to manually hit "Edit Revision" to edit it.

mlir/lib/Transforms/NormalizeMemRefs.cpp

179–181

Drop trivial braces.

[bondhugula]

mlir/include/mlir/Transforms/Passes.td

313

Normalize memrefs
(since there is no separate intra-procedural version).

mlir/include/mlir/Transforms/Utils.h

52

a non-dereferencing type -> a non-dereferencing user

75

handleDeallocOp -> replaceInDeallocOp ?

mlir/lib/Transforms/NormalizeMemRefs.cpp

84

Move the "TODO" to before the sentence.

94

-> isMemRefNormalizable?

Also, what's this mlir::detail::inLineOpResult ? Could you use the right type and not anything from mlir::detail. ArrayRef<Operation *>?

108

-> areMemRefsNormalizable?

130

Instead of getKind(), use isa.

mlir/test/Transforms/normalize-memrefs.mlir

201

Nit: ... because the function is returning the memref.

[bondhugula]

mlir/lib/Transforms/NormalizeMemRefs.cpp

110

Looks like you really don't need to collect/store the alloc ops in a vector. You can just run the check inline in the lambda.

[AlexEichenberger]

This comment was removed by AlexEichenberger.

[AlexEichenberger]

sorry, requested a revision by mistake. Repost as a comment below

[AlexEichenberger]

This patch does the advertised functionality well, as shown with the little example below:

#map0 = affine_map<(d0, d1) -> (d0, d1 floordiv 32, d1 mod 32)>
module {
  func @test(%in : memref<5x10xf32>, %out : memref<5x10xf32, #map0>) {
      affine.for %i = 0 to 5 {
          affine.for %j = 0 to 10 {
              %v = affine.load %in[%i, %j] :  memref<5x10xf32>
              affine.store %v, %out[%i, %j] : memref<5x10xf32, #map0>
          }
      }
      return
  }
  func @test_simplification() {
    %0 = alloc() : memref<5x10xf32>
    %1 = alloc() : memref<5x10xf32, #map0>
    //"test.test"(%0, %1) : (memref<5x10xf32>, memref<5x10xf32, #map0>) -> ()
    call @test(%0, %1) : (memref<5x10xf32>, memref<5x10xf32, #map0>) -> ()
    dealloc %1 : memref<5x10xf32, #map0>
    dealloc %0 : memref<5x10xf32>
    return
  }
}

transformed into

module {
  func @test(%arg0: memref<5x10xf32>, %arg1: memref<5x1x32xf32>) {
    affine.for %arg2 = 0 to 5 {
      affine.for %arg3 = 0 to 10 {
        %0 = affine.load %arg0[%arg2, %arg3] : memref<5x10xf32>
        affine.store %0, %arg1[%arg2, %arg3 floordiv 32, %arg3 mod 32] : memref<5x1x32xf32>
      }
    }
    return
  }
  func @test_simplification() {
    %0 = alloc() : memref<5x10xf32>
    %1 = alloc() : memref<5x1x32xf32>
    call @test(%0, %1) : (memref<5x10xf32>, memref<5x1x32xf32>) -> ()
    dealloc %1 : memref<5x1x32xf32>
    dealloc %0 : memref<5x10xf32>
    return
  }
}

where all the maps are eliminated from the load/store and alloc/dealloc.

However, a pattern that we see often is the lowering to external implementations not expressed in MLIR (think CUDNN calls or the like). In the above example, if we comment the "@test" lines and uncomment the line with "test.test", then the optimization will not remove any of the simplifications, as shown below.

func @test_simplification() {
  %0 = alloc() : memref<5x10xf32>
  %1 = alloc() : memref<5x10xf32, #map1>
  "test.test"(%0, %1) : (memref<5x10xf32>, memref<5x10xf32, #map1>) -> ()
  dealloc %1 : memref<5x10xf32, #map1>
  dealloc %0 : memref<5x10xf32>
  return
}

Is there a way to force the optimization through dialect operations? In our case, these maps were introduced especially to satisfy expected layout by the dialect "test." So we can take full responsibility that accesses within test.test will be fine. Happy with either a declarative approach or with a flag for a given dialect.

Thanks

[avarmapml]

Review changes and code formatting

mlir/lib/Transforms/NormalizeMemRefs.cpp

94

.getUsers() is returning Value::user_range and there is no known conversion from that to `ArrayRef<Operation *>.
I'm using Value::user_range as the type instead.

[avarmapml]

@AlexEichenberger : Such operation handling can be a part of subsequent revision.
Thank you.

[bondhugula]

mlir/lib/Transforms/NormalizeMemRefs.cpp

145–147

Can drop trivial braces.

179–181

Can drop trivial braces.

mlir/lib/Transforms/Utils/Utils.cpp

302

Nit: double dots at the end. Drop space before colon.

[avarmapml]

Minor formatting of code comment

[bondhugula]

Just a few minor comments.

mlir/include/mlir/Transforms/Utils.h

75

and we provide replacement -> where we replace

mlir/lib/Transforms/NormalizeMemRefs.cpp

26

"interprocedural memrefs" isn't meaningful - instead, memrefs passed across functions.

44

Nit: You can drop the module { } wrapping around.

[avarmapml]

Handled minor comments

[AlexEichenberger]

In D84490#2179268, @avarmapml wrote:

@AlexEichenberger : Such operation handling can be a part of subsequent revision.
Thank you.

Thanks, I think this would help quite a few folks, thanks. Maybe adding an interface that enables operations to trigger this behavior?

[bondhugula]

Can you please rebase and fix merge conflicts?

[avarmapml]

Resolved merge conflicts and involved another test case

[avarmapml]

Rebased to latest master branch

[avarmapml]

Clang-format error fixed

mlir/include/mlir/Transforms/Passes.h

 //===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 #include <functional>
 #include <limits>
 
 namespace mlir {
 
 class AffineForOp;
 class FuncOp;
 class ModuleOp;
 class Pass;
-template <typename T> class OperationPass;
+template <typename T>
+class OperationPass;
 
 /// Creates an instance of the BufferPlacement pass.
 std::unique_ptr<Pass> createBufferPlacementPass();
 
 /// Creates an instance of the Canonicalizer pass.
 std::unique_ptr<Pass> createCanonicalizerPass();
 
 /// Create a pass that removes unnecessary Copy operations.

 
 /// Creates a pass which performs sparse conditional constant propagation over
 /// nested operations.
 std::unique_ptr<Pass> createSCCPPass();
 
 /// Creates a pass which delete symbol operations that are unreachable. This
 /// pass may *only* be scheduled on an operation that defines a SymbolTable.
 std::unique_ptr<Pass> createSymbolDCEPass();
+
+/// Creates an interprocedural pass to normalize memrefs to have a trivial
+/// (identity) layout map.
+std::unique_ptr<OperationPass<ModuleOp>> createNormalizeMemRefsPass();
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_PASSES_H

mlir/include/mlir/Transforms/Passes.td

module {
       }
     }
     ```
   }];
   let constructor = "mlir::createMemRefDataFlowOptPass()";
 }
 
 def ParallelLoopCollapsing : Pass<"parallel-loop-collapsing"> {
   let summary = "Collapse parallel loops to use less induction variables";
   let constructor = "mlir::createParallelLoopCollapsingPass()";
   let options = [
     ListOption<"clCollapsedIndices0", "collapsed-indices-0", "unsigned",
                "Which loop indices to combine 0th loop index",
                "llvm::cl::MiscFlags::CommaSeparated">,
     ListOption<"clCollapsedIndices1", "collapsed-indices-1", "unsigned",
                "Which loop indices to combine into the position 1 loop index",
                "llvm::cl::MiscFlags::CommaSeparated">,
let description = [{
     ```
 
     See [Symbols and SymbolTables](SymbolsAndSymbolTables.md) for more
     information on `Symbols`.
   }];
   let constructor = "mlir::createSymbolDCEPass()";
 }
 
+def NormalizeMemRefs : Pass<"normalize-memrefs", "ModuleOp"> {
+  let summary = "Interprocedural normalize memrefs";
+  let constructor = "mlir::createNormalizeMemRefsPass()";
+}
 #endif // MLIR_TRANSFORMS_PASSES

mlir/include/mlir/Transforms/Utils.h

 //===- Utils.h - General transformation utilities ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 /// of its input list. `indexRemap`'s dimensional inputs are expected to
 /// correspond to memref's indices, and its symbolic inputs if any should be
 /// provided in `symbolOperands`.
 ///
 /// `domInstFilter`, if non-null, restricts the replacement to only those
 /// operations that are dominated by the former; similarly, `postDomInstFilter`
 /// restricts replacement to only those operations that are postdominated by it.
 ///
+/// 'allowNonDereferencingOps', if set, allows replacement of non-dereferencing 
+/// uses of a memref without any requirement for access index rewrites. The 
+/// default value of this flag variable is false.
+///
 /// Returns true on success and false if the replacement is not possible,
-/// whenever a memref is used as an operand in a non-dereferencing context,
-/// except for dealloc's on the memref which are left untouched. See comments at
-/// function definition for an example.
+/// whenever a memref is used as an operand in a non-dereferencing context and
+/// 'allowNonDereferencingOps' is false, except for dealloc's on the memref which
+/// are left untouched. See comments at function definition for an example.
 //
 //  Ex: to replace load %A[%i, %j] with load %Abuf[%t mod 2, %ii - %i, %j]:
 //  The SSA value corresponding to '%t mod 2' should be in 'extraIndices', and
 //  index remap will perform (%i, %j) -> (%ii - %i, %j), i.e., indexRemap = (d0,
 //  d1, d2) -> (d0 - d1, d2), and %ii will be the extra operand. Without any
 //  extra operands, note that 'indexRemap' would just be applied to existing
 //  indices (%i, %j).
 //  TODO: allow extraIndices to be added at any position.
 LogicalResult replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
                                        ArrayRef<Value> extraIndices = {},
                                        AffineMap indexRemap = AffineMap(),
                                        ArrayRef<Value> extraOperands = {},
                                        ArrayRef<Value> symbolOperands = {},
                                        Operation *domInstFilter = nullptr,
-                                       Operation *postDomInstFilter = nullptr);
+                                       Operation *postDomInstFilter = nullptr,
+                                       bool allowNonDereferencingOps = false);
 
 /// Performs the same replacement as the other version above but only for the
-/// dereferencing uses of `oldMemRef` in `op`.
+/// dereferencing uses of `oldMemRef` in `op`, except in cases where
+/// 'allowNonDereferencingOps' is set to true and we provide replacement for the
+/// non-dereferencing uses as well.
 LogicalResult replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
                                        Operation *op,
                                        ArrayRef<Value> extraIndices = {},
                                        AffineMap indexRemap = AffineMap(),
                                        ArrayRef<Value> extraOperands = {},
-                                       ArrayRef<Value> symbolOperands = {});
+                                       ArrayRef<Value> symbolOperands = {},
+                                       bool allowNonDereferencingOps = false);
 
 /// Rewrites the memref defined by this alloc op to have an identity layout map
 /// and updates all its indexing uses. Returns failure if any of its uses
 /// escape (while leaving the IR in a valid state).
 LogicalResult normalizeMemRef(AllocOp op);
 
+/// Computes the new memref type after trivializing the old memref map layout.
+/// It returns the old memref in case the trivializing was not needed or
+/// there was some failure while trivializing the map layout.
+MemRefType getNewMemRefType(MemRefType memrefType, OpBuilder builder, unsigned numSymbolicOperands);
+
 /// Creates and inserts into 'builder' a new AffineApplyOp, with the number of
 /// its results equal to the number of operands, as a composition
 /// of all other AffineApplyOps reachable from input parameter 'operands'. If
 /// different operands were drawing results from multiple affine apply ops,
 /// these will also be collected into a single (multi-result) affine apply op.
 /// The final results of the composed AffineApplyOp are returned in output
 /// parameter 'results'. Returns the affine apply op created.
 Operation *createComposedAffineApplyOp(OpBuilder &builder, Location loc,

mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp

 //===- SimplifyAffineStructures.cpp ---------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
for (auto attr : op->getAttrs()) {
         simplifyAndUpdateAttribute(op, attr.first, setAttr);
     }
 
     // The simplification of the attribute will likely simplify the op. Try to
     // fold / apply canonicalization patterns when we have affine dialect ops.
     if (isa<AffineForOp, AffineIfOp, AffineApplyOp>(op))
       applyOpPatternsAndFold(op, patterns);
   });
-
-  // Turn memrefs' non-identity layouts maps into ones with identity. Collect
-  // alloc ops first and then process since normalizeMemRef replaces/erases ops
-  // during memref rewriting.
-  SmallVector<AllocOp, 4> allocOps;
-  func.walk([&](AllocOp op) { allocOps.push_back(op); });
-  for (auto allocOp : allocOps) {
-    normalizeMemRef(allocOp);
-  }
 }

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   BufferPlacement.cpp
   Canonicalizer.cpp
   CopyRemoval.cpp
   CSE.cpp
   DialectConversion.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopInvariantCodeMotion.cpp
   MemRefDataFlowOpt.cpp
+  NormalizeMemRefs.cpp
   OpStats.cpp
   ParallelLoopCollapsing.cpp
   PipelineDataTransfer.cpp
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
   ViewOpGraph.cpp
   ViewRegionGraph.cpp

mlir/lib/Transforms/NormalizeMemRefs.cpp

+//===- NormalizeMemRefs.cpp ---------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements an interprocedural pass to normalize memref layouts.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Analysis/Utils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/FunctionSupport.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/Utils.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/Format.h"
+#include "llvm/Support/raw_ostream.h"
+
+#define DEBUG_TYPE "normalize-memrefs"
+
+using namespace mlir;
+
+/// All interprocedural memrefs with non-trivial layout maps are converted to
+/// ones with trivial identity layout ones.
+
+// Input :-
+// #tile = affine_map<(i) -> (i floordiv 4, i mod 4)>
+// func @matmul(%A: memref<16xf64, #tile>, %B: index, %C: memref<16xf64>) ->
+// (memref<16xf64, #tile>) {
+//   affine.for %arg3 = 0 to 16 {
+//         %a = affine.load %A[%arg3] : memref<16xf64, #tile>
+//         %p = mulf %a, %a : f64
+//         affine.store %p, %A[%arg3] : memref<16xf64, #tile>
+//   }
+//   %c = alloc() : memref<16xf64, #tile>
+//   %d = affine.load %c[0] : memref<16xf64, #tile>
+//   return %A: memref<16xf64, #tile>
+// }
+
+// Output :-
+// module {
+//   func @matmul(%arg0: memref<4x4xf64>, %arg1: index, %arg2: memref<16xf64>)
+//   -> memref<4x4xf64> {
+//     affine.for %arg3 = 0 to 16 {
+//       %2 = affine.load %arg0[%arg3 floordiv 4, %arg3 mod 4] : memref<4x4xf64>
+//       %3 = mulf %2, %2 : f64
+//       affine.store %3, %arg0[%arg3 floordiv 4, %arg3 mod 4] : memref<4x4xf64>
+//     }
+//     %0 = alloc() : memref<16xf64, #map0>
+//     %1 = affine.load %0[0] : memref<16xf64, #map0>
+//     return %arg0 : memref<4x4xf64>
+//   }
+// }
+
+namespace {
+
+struct NormalizeMemRefs : public NormalizeMemRefsBase<NormalizeMemRefs> {
+  void runOnOperation() override;
+  void runOnFunction(FuncOp, ModuleOp);
+};
+
+} // end anonymous namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> mlir::createNormalizeMemRefsPass() {
+  return std::make_unique<NormalizeMemRefs>();
+}
+
+void NormalizeMemRefs::runOnOperation() {
+  // Here we get hold of the module/operation that basically contains one
+  // region.
+  ModuleOp moduleOp = getOperation();
+
+  // We traverse each function within the module in order to normalize the
+  // memref type arguments.
+  // TODO(avarmapml): Handle external functions.
+  moduleOp.walk([&](FuncOp funcOp) { runOnFunction(funcOp, moduleOp); });
+}
+
+void NormalizeMemRefs::runOnFunction(FuncOp funcOp, ModuleOp moduleOp) {
+  // Turn memrefs' non-identity layouts maps into ones with identity. Collect
+  // alloc ops first and then process since normalizeMemRef replaces/erases ops
+  // during memref rewriting.
+  SmallVector<AllocOp, 4> allocOps;
+  funcOp.walk([&](AllocOp op) { allocOps.push_back(op); });
+  for (AllocOp allocOp : allocOps) {
+    normalizeMemRef(allocOp);
+  }
+  // We use this OpBuilder to create new memref layout later.
+  OpBuilder b(funcOp);
+
+  // Getting hold of the signature of the function before normalizing the
+  // inputs.
+  FunctionType ft = funcOp.getType();
+  SmallVector<Type, 8> argTypes;
+  SmallVector<Type, 4> resultTypes;
+
+  // Populating results with function's initial result (type) as
+  // this will be used to modify and set function's signature later.
+  for (unsigned retIndex : llvm::seq<unsigned>(0, funcOp.getNumResults())) {
+    resultTypes.push_back(ft.getResult(retIndex));
+  }
+
+  // Walk over each argument of a function to perform memref normalization (if
+  // any).
+  for (unsigned argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
+    Type argType = funcOp.getArgument(argIndex).getType();
+    MemRefType memrefType = argType.dyn_cast<MemRefType>();
+    // Check whether argument is of MemRef type.
+    if (!memrefType) {
+      // Any other argument type can simply be part of the final function
+      // signature.
+      argTypes.push_back(argType);
+      continue;
+    }
+    
+    // Fetch a new memref type after trivializing the old memref map layout.
+    MemRefType newMemRefType = getNewMemRefType(memrefType, b,
+                                                /*numSymbolicOperands=*/0);
+    if(newMemRefType == memrefType) {
+      // Either memrefType map layout was already trivial or the map couldn't
+      // be normalized.
+      argTypes.push_back(argType);
+      continue;
+    }
+    // Insert a new temporary argument with the new memref type.
+    funcOp.front().insertArgument(argIndex, newMemRefType);
+    BlockArgument newMemRef = funcOp.getArgument(argIndex);
+    BlockArgument oldMemRef = funcOp.getArgument(argIndex + 1);
+    AffineMap layoutMap = memrefType.getAffineMaps().front();
+    // Replace all uses of the old memref.
+    if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newMemRef,
+                                        /*extraIndices=*/{},
+                                        /*indexRemap=*/layoutMap,
+                                        /*extraOperands=*/{},
+                                        /*symbolOperands=*/{},
+                                        /*domInstFilter=*/nullptr,
+                                        /*postDomInstFilter=*/nullptr,
+                                        /*allowNonDereferencingOps=*/true))) {
+      // If it failed (due to escapes for example), bail out. Removing the
+      // temporary argument inserted previously.
+      funcOp.front().eraseArgument(argIndex);
+      continue;
+    }
+
+    // Since in this pass the objective is to normalize the layout maps of
+    // the memref arguments and replace the uses accordingly, we
+    // check if the function return type uses the same old memref type.
+    // TODO(avarmapml): Check - A function's return type might have a
+    // different memref layout and a map.
+    for (unsigned retIndex : llvm::seq<unsigned>(0, funcOp.getNumResults())) {
+      if (resultTypes[retIndex] == memrefType)
+        resultTypes[retIndex] = newMemRef.getType();
+    }
+
+    // All uses for the argument with old memref type were replaced
+    // successfully. So we remove the old argument now.
+    // TODO(avarmapml): replaceAllUsesWith.
+    funcOp.front().eraseArgument(argIndex + 1);
+
+    // Add the new type to the function signature later.
+    argTypes.push_back(static_cast<Type>(newMemRefType));
+  }
+
+  ArrayRef<Type> inputs = makeArrayRef(argTypes);
+  ArrayRef<Type> results = makeArrayRef(resultTypes);
+  // We create a new function type and modify the function signature with this
+  // new type.
+  FunctionType newFuncType = FunctionType::get(/*inputs=*/inputs,
+                                               /*results=*/results,
+                                               /**context=*/&getContext());
+
+  // We get hold of all the symbolic uses of this function within the given
+  // region.
+  Optional<SymbolTable::UseRange> symbolUses = SymbolTable::getSymbolUses(
+      funcOp.getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
+          .getValue(),
+      &(moduleOp.getBodyRegion()));
+  // We iterate over all symbolic uses of the function and update the return
+  // type at the caller site.
+  for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
+    Operation *callOp = symbolUse.getUser();
+    // We build a new CallOp to reflect the updated return type of the function.
+    OpBuilder builder(callOp);
+    OperationState state(callOp->getLoc(), callOp->getName());
+    state.operands.reserve(callOp->getNumOperands());
+    state.operands.append(callOp->operand_begin(), callOp->operand_end());
+    // We take care only of the return type from the function signature because
+    // the operand type will get changed later at the caller function.
+    state.addTypes(results);
+    StringRef callee = static_cast<CallOp>(callOp).getCallee();
+    state.addAttribute("callee", builder.getSymbolRefAttr(callee));
+    Operation *newCallOp = builder.createOperation(state);
+    bool replacingMemRefUsesFailed = false;
+    // A function might return more than one result. So here we loop over
+    // all the results returned and replace uses of any result of type Memref
+    // whose map layout has changed.
+    for (unsigned resIndex : llvm::seq<unsigned>(0, callOp->getNumResults())) {
+      OpResult oldMemRef = callOp->getResult(resIndex);
+      OpResult newMemRef = newCallOp->getResult(resIndex);
+      // If the result type is not a Memref or if there is no change in the
+      // Memref map layout, there is no need to call replaceAllMemRefUsesWith.
+      if (oldMemRef.getType() == newMemRef.getType())
+        continue;
+      AffineMap layoutMap =
+          oldMemRef.getType().dyn_cast<MemRefType>().getAffineMaps().front();
+      if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newMemRef,
+                                          /*extraIndices=*/{},
+                                          /*indexRemap=*/layoutMap,
+                                          /*extraOperands=*/{},
+                                          /*symbolOperands=*/{},
+                                          /*domInstFilter=*/nullptr,
+                                          /*postDomInstFilter=*/nullptr,
+                                          /*allowDereferencingOps=*/true))) {
+        // If it failed (due to escapes for example), bail out.
+        newCallOp->erase();
+        replacingMemRefUsesFailed = true;
+        break;
+      }
+    }
+    if (replacingMemRefUsesFailed)
+      continue;
+    callOp->replaceAllUsesWith(newCallOp);
+    callOp->erase();
+  }
+  funcOp.setType(newFuncType);
+}

mlir/lib/Transforms/Utils/Utils.cpp

 //===- Utils.cpp ---- Misc utilities for code and data transformation -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 }
 
 // Perform the replacement in `op`.
 LogicalResult mlir::replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
                                              Operation *op,
                                              ArrayRef<Value> extraIndices,
                                              AffineMap indexRemap,
                                              ArrayRef<Value> extraOperands,
-                                             ArrayRef<Value> symbolOperands) {
+                                             ArrayRef<Value> symbolOperands,
+                                             bool allowNonDereferencingOps) {
   unsigned newMemRefRank = newMemRef.getType().cast<MemRefType>().getRank();
   (void)newMemRefRank; // unused in opt mode
   unsigned oldMemRefRank = oldMemRef.getType().cast<MemRefType>().getRank();
   (void)oldMemRefRank; // unused in opt mode
   if (indexRemap) {
     assert(indexRemap.getNumSymbols() == symbolOperands.size() &&
            "symbolic operand count mismatch");
     assert(indexRemap.getNumInputs() ==
            extraOperands.size() + oldMemRefRank + symbolOperands.size());
     assert(indexRemap.getNumResults() + extraIndices.size() == newMemRefRank);
   } else {
     assert(oldMemRefRank + extraIndices.size() == newMemRefRank);
   }
 
   // Assert same elemental type.
   assert(oldMemRef.getType().cast<MemRefType>().getElementType() ==
          newMemRef.getType().cast<MemRefType>().getElementType());
 
-  if (!isMemRefDereferencingOp(*op))
-    // Failure: memref used in a non-dereferencing context (potentially
-    // escapes); no replacement in these cases.
-    return failure();
-
   SmallVector<unsigned, 2> usePositions;
   for (const auto &opEntry : llvm::enumerate(op->getOperands())) {
     if (opEntry.value() == oldMemRef)
       usePositions.push_back(opEntry.index());
   }
 
   // If memref doesn't appear, nothing to do.
   if (usePositions.empty())
     return success();
 
   if (usePositions.size() > 1) {
     // TODO: extend it for this case when needed (rare).
     assert(false && "multiple dereferencing uses in a single op not supported");
     return failure();
   }
 
   unsigned memRefOperandPos = usePositions.front();
 
   OpBuilder builder(op);
+  OperationState state(op->getLoc(), op->getName());
+  // The following checks if op is dereferencing memref and performs the access
+  // index rewrites.
+  if (!isMemRefDereferencingOp(*op)) {
+    if (!allowNonDereferencingOps)
+      // Failure: memref used in a non-dereferencing context (potentially
+      // escapes); no replacement in these cases unless allowNonDereferencingOps
+      // is set.
+      return failure();
+    // For non-dereferencing op we simply replace the memref type.
+    state.operands.reserve(op->getNumOperands() + extraIndices.size());
+    // Insert the operands ahead of the oldMemRef type in argument list.
+    state.operands.append(op->operand_begin(),
+                          op->operand_begin() + memRefOperandPos);
+    // Insert the new memref value.
+    state.operands.push_back(newMemRef);
+    // Insert other operands that follow the oldMemRef type in argumemt list.
+    state.operands.append(op->operand_begin() + memRefOperandPos + 1,
+                          op->operand_end());
+    state.types.reserve(op->getNumResults());
+    for (auto result : op->getResults())
+      state.types.push_back(result.getType());
+
+    // For CallOps we need to add an attribute 'callee' whose value will be the
+    // function name.
+    if (isa<CallOp>(*op)) {
+      StringRef callee = static_cast<CallOp>(op).getCallee();
+      state.addAttribute("callee", builder.getSymbolRefAttr(callee));
+    }
+
+    // Create the new operation.
+    Operation *repOp = builder.createOperation(state);
+    op->replaceAllUsesWith(repOp);
+    op->erase();
+
+    return success();
+  }
+  // Perform index rewrites for the dereferencing op and then replace the op
   NamedAttribute oldMapAttrPair = getAffineMapAttrForMemRef(op, oldMemRef);
   AffineMap oldMap = oldMapAttrPair.second.cast<AffineMapAttr>().getValue();
   unsigned oldMapNumInputs = oldMap.getNumInputs();
   SmallVector<Value, 4> oldMapOperands(
       op->operand_begin() + memRefOperandPos + 1,
       op->operand_begin() + memRefOperandPos + 1 + oldMapNumInputs);
 
   // Apply 'oldMemRefOperands = oldMap(oldMapOperands)'.
   SmallVector<Value, 4> oldMemRefOperands;
   SmallVector<Value, 4> affineApplyOps;
   oldMemRefOperands.reserve(oldMemRefRank);
   if (oldMap != builder.getMultiDimIdentityMap(oldMap.getNumDims())) {
     for (auto resultExpr : oldMap.getResults()) {
       auto singleResMap = AffineMap::get(oldMap.getNumDims(),
                                          oldMap.getNumSymbols(), resultExpr);
       auto afOp = builder.create<AffineApplyOp>(op->getLoc(), singleResMap,
                                                 oldMapOperands);
       oldMemRefOperands.push_back(afOp);
       affineApplyOps.push_back(afOp);
     }
   } else {
-    oldMemRefOperands.append(oldMapOperands.begin(), oldMapOperands.end());
+    oldMemRefOperands.assign(oldMapOperands.begin(), oldMapOperands.end());
   }
 
   // Construct new indices as a remap of the old ones if a remapping has been
   // provided. The indices of a memref come right after it, i.e.,
   // at position memRefOperandPos + 1.
   SmallVector<Value, 4> remapOperands;
   remapOperands.reserve(extraOperands.size() + oldMemRefRank +
                         symbolOperands.size());
for (auto resultExpr : indexRemap.getResults()) {
           indexRemap.getNumDims(), indexRemap.getNumSymbols(), resultExpr);
       auto afOp = builder.create<AffineApplyOp>(op->getLoc(), singleResMap,
                                                 remapOperands);
       remapOutputs.push_back(afOp);
       affineApplyOps.push_back(afOp);
     }
   } else {
     // No remapping specified.
-    remapOutputs.append(remapOperands.begin(), remapOperands.end());
+    remapOutputs.assign(remapOperands.begin(), remapOperands.end());
   }
 
   SmallVector<Value, 4> newMapOperands;
   newMapOperands.reserve(newMemRefRank);
 
   // Prepend 'extraIndices' in 'newMapOperands'.
-  for (auto extraIndex : extraIndices) {
+  for (Value extraIndex : extraIndices) {
     assert(extraIndex.getDefiningOp()->getNumResults() == 1 &&
            "single result op's expected to generate these indices");
     assert((isValidDim(extraIndex) || isValidSymbol(extraIndex)) &&
            "invalid memory op index");
     newMapOperands.push_back(extraIndex);
   }
 
   // Append 'remapOutputs' to 'newMapOperands'.
   newMapOperands.append(remapOutputs.begin(), remapOutputs.end());
 
   // Create new fully composed AffineMap for new op to be created.
   assert(newMapOperands.size() == newMemRefRank);
   auto newMap = builder.getMultiDimIdentityMap(newMemRefRank);
   // TODO: Avoid creating/deleting temporary AffineApplyOps here.
   fullyComposeAffineMapAndOperands(&newMap, &newMapOperands);
   newMap = simplifyAffineMap(newMap);
   canonicalizeMapAndOperands(&newMap, &newMapOperands);
   // Remove any affine.apply's that became dead as a result of composition.
-  for (auto value : affineApplyOps)
+  for (Value value : affineApplyOps)
     if (value.use_empty())
       value.getDefiningOp()->erase();
 
   // Construct the new operation using this memref.
-  OperationState state(op->getLoc(), op->getName());
   state.operands.reserve(op->getNumOperands() + extraIndices.size());
   // Insert the non-memref operands.
   state.operands.append(op->operand_begin(),
                         op->operand_begin() + memRefOperandPos);
   // Insert the new memref value.
   state.operands.push_back(newMemRef);
 
   // Insert the new memref map operands.
   state.operands.append(newMapOperands.begin(), newMapOperands.end());
 
   // Insert the remaining operands unmodified.
   state.operands.append(op->operand_begin() + memRefOperandPos + 1 +
                             oldMapNumInputs,
                         op->operand_end());
 
   // Result types don't change. Both memref's are of the same elemental type.
   state.types.reserve(op->getNumResults());
   for (auto result : op->getResults())
     state.types.push_back(result.getType());
 
   // Add attribute for 'newMap', other Attributes do not change.
   auto newMapAttr = AffineMapAttr::get(newMap);
   for (auto namedAttr : op->getAttrs()) {
-    if (namedAttr.first == oldMapAttrPair.first) {
+    if (namedAttr.first == oldMapAttrPair.first)
       state.attributes.push_back({namedAttr.first, newMapAttr});
-    } else {
+    else
       state.attributes.push_back(namedAttr);
-    }
-  }
 
   // Create the new operation.
   auto *repOp = builder.createOperation(state);
   op->replaceAllUsesWith(repOp);
   op->erase();
 
   return success();
 }
 
-LogicalResult mlir::replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
-                                             ArrayRef<Value> extraIndices,
-                                             AffineMap indexRemap,
-                                             ArrayRef<Value> extraOperands,
-                                             ArrayRef<Value> symbolOperands,
-                                             Operation *domInstFilter,
-                                             Operation *postDomInstFilter) {
+LogicalResult mlir::replaceAllMemRefUsesWith(
+    Value oldMemRef, Value newMemRef, ArrayRef<Value> extraIndices,
+    AffineMap indexRemap, ArrayRef<Value> extraOperands,
+    ArrayRef<Value> symbolOperands, Operation *domInstFilter,
+    Operation *postDomInstFilter, bool allowNonDereferencingOps) {
   unsigned newMemRefRank = newMemRef.getType().cast<MemRefType>().getRank();
   (void)newMemRefRank; // unused in opt mode
   unsigned oldMemRefRank = oldMemRef.getType().cast<MemRefType>().getRank();
   (void)oldMemRefRank;
   if (indexRemap) {
     assert(indexRemap.getNumSymbols() == symbolOperands.size() &&
            "symbol operand count mismatch");
     assert(indexRemap.getNumInputs() ==
for (auto *op : oldMemRef.getUsers()) {
     // Skip dealloc's - no replacement is necessary, and a memref replacement
     // at other uses doesn't hurt these dealloc's.
     if (isa<DeallocOp>(op))
       continue;
 
     // Check if the memref was used in a non-dereferencing context. It is fine
     // for the memref to be used in a non-dereferencing way outside of the
     // region where this replacement is happening.
-    if (!isMemRefDereferencingOp(*op))
-      // Failure: memref used in a non-dereferencing op (potentially escapes);
-      // no replacement in these cases.
+    if (!isMemRefDereferencingOp(*op)) {
+      // Currently we support the following non-dereferencing types to be a
+      // candidate for replacement : DeallocOp, CallOp and ReturnOp.
+      // TODO: Add support for other kinds of ops.
+      if (!allowNonDereferencingOps)
+        return failure();
+      if (!(isa<CallOp, ReturnOp>(*op)))
         return failure();
+    }
 
     // We'll first collect and then replace --- since replacement erases the op
     // that has the use, and that op could be postDomFilter or domFilter itself!
     opsToReplace.insert(op);
   }
 
   for (auto *op : opsToReplace) {
-    if (failed(replaceAllMemRefUsesWith(oldMemRef, newMemRef, op, extraIndices,
-                                        indexRemap, extraOperands,
-                                        symbolOperands)))
+    if (failed(replaceAllMemRefUsesWith(
+            oldMemRef, newMemRef, op, extraIndices, indexRemap, extraOperands,
+            symbolOperands, allowNonDereferencingOps)))
       llvm_unreachable("memref replacement guaranteed to succeed here");
   }
 
   return success();
 }
 
 /// Given an operation, inserts one or more single result affine
 /// apply operations, results of which are exclusively used by this operation
void mlir::createAffineComputationSlice(
   for (unsigned idx = 0, e = newOperands.size(); idx < e; idx++) {
     opInst->setOperand(idx, newOperands[idx]);
   }
 }
 
 // TODO: Currently works for static memrefs with a single layout map.
 LogicalResult mlir::normalizeMemRef(AllocOp allocOp) {
   MemRefType memrefType = allocOp.getType();
-  unsigned rank = memrefType.getRank();
-  if (rank == 0)
-    return success();
-
-  auto layoutMaps = memrefType.getAffineMaps();
   OpBuilder b(allocOp);
-  if (layoutMaps.size() != 1)
+
+  // Fetch a new memref type after trivializing the old memref map layout.
+  MemRefType newMemRefType = getNewMemRefType(memrefType, b, allocOp.getNumSymbolicOperands());
+  if(newMemRefType == memrefType)
+    // Either memrefType map layout was already trivial or the map couldn't
+    // be normalized.
     return failure();
 
-  AffineMap layoutMap = layoutMaps.front();
+  Value oldMemRef = allocOp.getResult();
+  // Check whether all the uses of oldMemRef is either dereferencing uses
+  // or the op is of type : DeallocOp, CallOp or a ReturnOp. Only if these
+  // constraints are satisfied will the op become candidate for replacement.
+  if (llvm::any_of(oldMemRef.getUsers(), [](Operation *op) {
+        if (isMemRefDereferencingOp(*op))
+          return false;
+        return (!isa<DeallocOp, CallOp>(*op));
+      }))
+    return failure();
 
-  // Nothing to do for identity layout maps.
-  if (layoutMap == b.getMultiDimIdentityMap(rank))
-    return success();
+  SmallVector<Value, 4> symbolOperands(allocOp.getSymbolicOperands());
+  AllocOp newAlloc = b.create<AllocOp>(allocOp.getLoc(), newMemRefType);
+  AffineMap layoutMap = memrefType.getAffineMaps().front();
+  // Replace all uses of the old memref.
+  if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
+                                      /*extraIndices=*/{},
+                                      /*indexRemap=*/layoutMap,
+                                      /*extraOperands=*/{},
+                                      /*symbolOperands=*/symbolOperands,
+                                      /*domInstFilter=*/nullptr,
+                                      /*postDomInstFilter=*/nullptr,
+                                      /*allowDereferencingOps=*/true))) {
+    // If it failed (due to escapes for example), bail out.
+    newAlloc.erase();
+    return failure();
+  }
+  // Replace any uses of the original alloc op and erase it. All remaining uses
+  // have to be dealloc's; RAMUW above would've failed otherwise.
+  assert(llvm::all_of(oldMemRef.getUsers(),
+                      [](Operation *op) { return isa<DeallocOp>(op); }));
+  oldMemRef.replaceAllUsesWith(newAlloc);
+  allocOp.erase();
+}
+
+MemRefType mlir::getNewMemRefType(MemRefType memrefType, OpBuilder b, unsigned numSymbolicOperands) {
+  unsigned rank = memrefType.getRank();
+  if (rank == 0)
+    return memrefType;
+
+  ArrayRef<AffineMap> layoutMaps = memrefType.getAffineMaps();
+  if (layoutMaps.empty() ||
+      layoutMaps.front() == b.getMultiDimIdentityMap(rank)) {
+    // Either no maps is associated with this memref or this memref has
+    // a trivial (identity) map.
+    return memrefType;
+  }
 
   // We don't do any checks for one-to-one'ness; we assume that it is
   // one-to-one.
 
   // TODO: Only for static memref's for now.
   if (memrefType.getNumDynamicDims() > 0)
-    return failure();
+    return memrefType;
 
-  // We have a single map that is not an identity map. Create a new memref with
-  // the right shape and an identity layout map.
-  auto shape = memrefType.getShape();
-  FlatAffineConstraints fac(rank, allocOp.getNumSymbolicOperands());
+  // We have a single map that is not an identity map. Create a new memref
+  // with the right shape and an identity layout map.
+  ArrayRef<int64_t> shape = memrefType.getShape();
+  // FlatAffineConstraint may later on use symbolicOperands.
+  FlatAffineConstraints fac(rank, numSymbolicOperands);
   for (unsigned d = 0; d < rank; ++d) {
     fac.addConstantLowerBound(d, 0);
     fac.addConstantUpperBound(d, shape[d] - 1);
   }
-
-  // We compose this map with the original index (logical) space to derive the
-  // upper bounds for the new index space.
+  // We compose this map with the original index (logical) space to derive
+  // the upper bounds for the new index space.
+  AffineMap layoutMap = layoutMaps.front();
   unsigned newRank = layoutMap.getNumResults();
   if (failed(fac.composeMatchingMap(layoutMap)))
-    // TODO: semi-affine maps.
-    return failure();
-
+    return memrefType;
+  // TODO: Handle semi-affine maps.
   // Project out the old data dimensions.
   fac.projectOut(newRank, fac.getNumIds() - newRank - fac.getNumLocalIds());
   SmallVector<int64_t, 4> newShape(newRank);
   for (unsigned d = 0; d < newRank; ++d) {
     // The lower bound for the shape is always zero.
     auto ubConst = fac.getConstantUpperBound(d);
-    // For a static memref and an affine map with no symbols, this is always
-    // bounded.
+    // For a static memref and an affine map with no symbols, this is
+    // always bounded.
     assert(ubConst.hasValue() && "should always have an upper bound");
     if (ubConst.getValue() < 0)
       // This is due to an invalid map that maps to a negative space.
-      return failure();
+      return memrefType;
     newShape[d] = ubConst.getValue() + 1;
   }
 
-  auto oldMemRef = allocOp.getResult();
-  SmallVector<Value, 4> symbolOperands(allocOp.getSymbolicOperands());
-
+  // Create the new memref type after trivializing the old layout map.
   MemRefType newMemRefType =
       MemRefType::Builder(memrefType)
           .setShape(newShape)
           .setAffineMaps(b.getMultiDimIdentityMap(newRank));
-  auto newAlloc = b.create<AllocOp>(allocOp.getLoc(), newMemRefType);
 
-  // Replace all uses of the old memref.
-  if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
-                                      /*extraIndices=*/{},
-                                      /*indexRemap=*/layoutMap,
-                                      /*extraOperands=*/{},
-                                      /*symbolOperands=*/symbolOperands))) {
-    // If it failed (due to escapes for example), bail out.
-    newAlloc.erase();
-    return failure();
-  }
-  // Replace any uses of the original alloc op and erase it. All remaining uses
-  // have to be dealloc's; RAMUW above would've failed otherwise.
-  assert(llvm::all_of(oldMemRef.getUsers(),
-                      [](Operation *op) { return isa<DeallocOp>(op); }));
-  oldMemRef.replaceAllUsesWith(newAlloc);
-  allocOp.erase();
-  return success();
-}
+  return newMemRefType;
+}

mlir/test/Transforms/normalize-memrefs.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -simplify-affine-structures %s | FileCheck %s
+// RUN: mlir-opt -normalize-memrefs -allow-unregistered-dialect %s | FileCheck %s
+
+// This file tests whether the memref type having non-trivial map layouts
+// are normalized to trivial (identity) layouts.
 
 // CHECK-LABEL: func @permute()
 func @permute() {
   %A = alloc() : memref<64x256xf32, affine_map<(d0, d1) -> (d1, d0)>>
   affine.for %i = 0 to 64 {
     affine.for %j = 0 to 256 {
       %1 = affine.load %A[%i, %j] : memref<64x256xf32, affine_map<(d0, d1) -> (d1, d0)>>
       "prevent.dce"(%1) : (f32) -> ()
func @semi_affine_layout_map(%s0: index, %s1: index) {
   affine.for %i = 0 to 256 {
     affine.for %j = 0 to 1024 {
       // CHECK: memref<256x1024xf32, #map{{[0-9]+}}>
       affine.load %A[%i, %j] : memref<256x1024xf32, affine_map<(d0, d1)[s0, s1] -> (d0*s0 + d1*s1)>>
     }
   }
   return
 }
+
+#tile = affine_map<(i) -> (i floordiv 4, i mod 4)>
+
+// Following test cases check the inter-procedural memref normalization.
+
+// Test case 1: Checks whether memref map layout normalization takes place successfully for recursive function calls.
+// CHECK-LABEL: func @recursive
+// CHECK-SAME: (%[[A:arg[0-9]+]]: memref<4x4xf64>) -> memref<4x4xf64>
+func @recursive(%A: memref<16xf64, #tile>) -> (memref<16xf64, #tile>) {
+  %M = call @recursive(%A) : (memref<16xf64, #tile>) -> (memref<16xf64, #tile>)
+  %c = alloc() : memref<8xf64>
+  %b = constant 4.5 : f64
+  %R = call @argument_type(%A, %b, %c) : (memref<16xf64, #tile>, f64, memref<8xf64>) -> (memref<16xf64, #tile>)
+  return %M: memref<16xf64, #tile>
+}
+
+// CHECK: [[M:%[0-9]+]] = call @recursive(%[[A]]) : (memref<4x4xf64>) -> memref<4x4xf64>
+// CHECK: [[R:%[0-9]+]] = call @argument_type(%[[A]], %cst, %[[c:[0-9]+]]) : (memref<4x4xf64>, f64, memref<8xf64>) -> memref<4x4xf64>
+// CHECK: return [[M]] : memref<4x4xf64>
+
+
+// Test case 2: Checks whether function arguments of memref type having non-trivial map layout get normalized.
+//              Also checks interprocedural memref layout normalization
+// CHECK-LABEL: func @argument_type
+// CHECK-SAME: (%[[A:arg[0-9]+]]: memref<4x4xf64>, %[[B:arg[0-9]+]]: f64, %[[C:arg[0-9]+]]: memref<8xf64>) -> memref<4x4xf64>
+func @argument_type(%A: memref<16xf64, #tile>, %B: f64, %C: memref<8xf64>) -> (memref<16xf64, #tile>) {
+  %a = affine.load %A[0] : memref<16xf64, #tile>
+  %p = mulf %a, %a : f64
+  affine.store %p, %A[10] : memref<16xf64, #tile>
+  %M = call @recursive(%A) : (memref<16xf64, #tile>) -> (memref<16xf64, #tile>)
+  %m = affine.load %M[1] : memref<16xf64, #tile>
+  return %A: memref<16xf64, #tile>
+}
+
+// CHECK: %[[a:[0-9]+]] = affine.load %[[A]][0, 0] : memref<4x4xf64>
+// CHECK: %[[p:[0-9]+]] = mulf %[[a]], %[[a]] : f64
+// CHECK: affine.store %[[p]], %[[A]][2, 2] : memref<4x4xf64>
+// CHECK: [[MEM:%[0-9]+]] = call @recursive(%[[A]]) : (memref<4x4xf64>) -> memref<4x4xf64>
+// CHECK: %[[b:[0-9]+]] = affine.load [[MEM]][0, 1] : memref<4x4xf64>
+// CHECK: return %[[A]] : memref<4x4xf64>
+
+// Test case 3: Check normalization for multiple memrefs in a function argument list.
+// CHECK-LABEL: func @multiple_arg_normalize
+// CHECK-SAME: (%[[A:arg[0-9]+]]: memref<4x4xf64>, %[[B:arg[0-9]+]]: memref<2x4xf64>) -> memref<2x4xf64>
+func @multiple_arg_normalize(%A: memref<16xf64, #tile>, %B: memref<8xf64, #tile>) -> (memref<8xf64, #tile>) {
+  %a = affine.load %A[0] : memref<16xf64, #tile>
+  return %B: memref<8xf64, #tile>
+}
+
+// CHECK: %[[argA:[0-9]+]] = affine.load %[[A]][0, 0] : memref<4x4xf64>
+// CHECK: return %[[B]] : memref<2x4xf64>