[mlir] Normalize dynamic memrefs with a map of tiled-layout.

[imaihal]

Summary

Steps for normalizing dynamic memrefs for tiled layout map

1. Check if original map is tiled layout. Only tiled layout is supported.
2. Create normalized memrefType. Dimensions that include dynamic dimensions in the map output will be dynamic dimensions.
3. Create new maps to calculate each dimension size of new memref. In tiled layout, the dimension size can be calculated by replacing "floordiv <tile size>" with "ceildiv <tile size>" and "mod <tile size>" with "<tile size>".
4. Create AffineApplyOp to apply the new maps. The output of AffineApplyOp is dynamicSizes for new AllocOp.
5. Add the new dynamic sizes in new AllocOp.

This patch also set MemRefsNormalizable trant in CastOp and DimOp since
they used with dynamic memrefs.

[imaihal]

Clean up.

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

438–439

Please use empty().

No need of the else below.

462–465

Triple /// comments here and everywhere below.

[avarmapml]

A few minor comments.

mlir/lib/Transforms/Utils/Utils.cpp

455–458

No need of the else.

531

Terminate the comment with fullstop.

544

Terminate the comment with fullstop.

648–650

else block should also have the braces since the then block has it.

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

419

Doc comment please.

421

is a tiled ...

422

This is really a useful utility to have - thanks!

450

tileCount - please expand it out for readability.

[imaihal]

Rebased and Reflected comments.

[imaihal]

@bondhugula @avarmapml Thanks for your reviews! I reflected your comments.

mlir/lib/Transforms/Utils/Utils.cpp

419

Updated comments and moved them before createNewDynamicSizes()

[imaihal]

Rebased.

[imaihal]

@bondhugula @avarmapml Thanks for your review! I reflected your comments. Do you have any other comments on this?

[bondhugula]

Ping @avarmapml

[bondhugula]

Some style related comments.

mlir/lib/Transforms/NormalizeMemRefs.cpp

39 ↗

(On Diff #328164)

It may be unclear why this comment is specific to dynamic memrefs here. You can just say that the affine.apply op is needed to normalize dynamic memrefs.

mlir/lib/Transforms/Utils/Utils.cpp

382

Doc comment please explaining the meaning of this enum and its values.

396

Instead of getKind() - please use isa<AffineConsta.... Likewise above and anywhere else.

410

const std::pair... &ts to avoid a copy.

435

You don't need to pass a builder - this method is not constructing IR. You only need the context which can be passed.

Also, please move the output argument to the end (inMemrefTypeDynDims).

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

14–20

These would be hard to maintain and the tests below would be less readable with all definitions here. Instead, please use -split-input-file and the ----- separator. The maps will be relative/local in that case and can be defined right above the function.

93

of not tiled layout -> that are not tiled layout maps

[bondhugula]

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

30–31

I assume these dim ops will still be around after this pass but they would disappear after a -canonicalize if the alloc op they were used in is erased.

[imaihal]

Reflected comments.

[imaihal]

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

30–31

@bondhugula Thanks for your comments!
I checked generated code after canonicalization. In this case, dim ops are still existed, but some affine_maps are erased. The code before canonicalization shows what I implemented. So, don't I have to update this output? Or, should I use the output after canonicalization?

• Before canonicalization

#map0 = affine_map<(d0, d1, d2, d3) -> (d1)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d2 ceildiv 32)>
#map2 = affine_map<(d0, d1, d2, d3) -> (32)>
module  {
  func @test_norm_dynamic12(%arg0: memref<1x?x?x1x?x64xf32>) {
    %c1 = constant 1 : index
    %c2 = constant 2 : index
    %0 = dim %arg0, %c1 : memref<1x?x?x1x?x64xf32>
    %1 = dim %arg0, %c2 : memref<1x?x?x1x?x64xf32>
    %c1_0 = constant 1 : index
    %c14 = constant 14 : index
    %2 = affine.apply #map0(%c1_0, %0, %1, %c14)
    %3 = affine.apply #map1(%c1_0, %0, %1, %c14)
    %4 = affine.apply #map2(%c1_0, %0, %1, %c14)
    %5 = alloc(%2, %3, %4) : memref<1x?x?x1x?x64xf32>
    "test.op_norm"(%arg0, %5) : (memref<1x?x?x1x?x64xf32>, memref<1x?x?x1x?x64xf32>) -> ()
    dealloc %5 : memref<1x?x?x1x?x64xf32>
    return
  }
}

• After canonicalization

#map = affine_map<()[s0] -> (s0 ceildiv 32)>
module  {
  func @test_norm_dynamic12(%arg0: memref<1x?x?x1x?x64xf32>) {
    %c1 = constant 1 : index
    %c2 = constant 2 : index
    %0 = dim %arg0, %c1 : memref<1x?x?x1x?x64xf32>
    %1 = dim %arg0, %c2 : memref<1x?x?x1x?x64xf32>
    %2 = affine.apply #map()[%1]
    %3 = alloc(%0, %2) : memref<1x?x?x1x32x64xf32>
    %4 = memref_cast %3 : memref<1x?x?x1x32x64xf32> to memref<1x?x?x1x?x64xf32>
    "test.op_norm"(%arg0, %4) : (memref<1x?x?x1x?x64xf32>, memref<1x?x?x1x?x64xf32>) -> ()
    dealloc %3 : memref<1x?x?x1x32x64xf32>
    return
  }
}

[imaihal]

Removed canonicalize.

[imaihal]

@bondhugula @avarmapml Thanks for your review! I reflected your comments. Do you have additional comments?

[imaihal]

@bondhugula @avarmapml Sorry again, but I would like to use this patch in my use case. Could you merge this if you can approve?

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

394

A tiled layout map has to be clearly defined here so that it's clear what the formal condition to check for is. Please see comment further below.

420

While this handles the common case, this method isn't sound in general: because you'd have to check that for each d_i floordiv t_s, there exists one and only one d_i mod t_s and the former comes before the latter? But even this part isn't clear because a tiled layout map isn't defined in the doc comment of this method - so it'd be ultimately unclear as to what the formal condition to check here is.

427–429

Nit: return tileCount == tileSizes.size();

[imaihal]

Rebased and reflected comments.

[imaihal]

Fixed code that accidentally deleted a previous fix, and updated comments.

[imaihal]

Fixed a comment.

[imaihal]

@bondhugula Thanks for your review! I reflected your comments. Could you review again?

mlir/lib/Transforms/Utils/Utils.cpp

420

@bondhugula I used AffineExpr.walk here, but it was not correct. I removed it.

This isTiledLayoutMap works as follows in #map_tiled1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, (d2 floordiv 4) floordiv 32, (d3 mod 8) floordiv 64, (d2 floordiv 4) mod 32, (d3 mod 8) mod 64)>,

1. Finds "<AffineExpr> floordiv <AffineConstantExpr>".

In this example, they are (d2 floordiv 4) floordiv 32 and (d3 mod 8) floordiv 64. So, tileSizes includes {d0 floordiv 4, 32} and {d3 mod 8, 64}.

2. Finds "<AffineExpr> mod <AffineConstantExpr>", and compares them with the tileSizes

In this example, they are (d2 floordiv 4) mod 32 and (d3 mod 8) mod 64, and the pair of {d2 floordive 4, 32} and {d3 mod 8, 64} exists in the tileSize. So it returns true.

[bondhugula]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

246

Please capture this change in the commit summary as well - thanks.

396

Likewise - please add a brief line on this in the commit summary.

mlir/lib/Transforms/Utils/Utils.cpp

395

But this condition is technically wrong - it's a necessary one but no where close to the sufficient condition. The number of floordiv's being equal to the number of mod's for a dim expression doesn't make the affine map a tiled layout map. It needs to have a specific structure.

[imaihal]

Updated condition of tiled-layout map used in isTiledLayoutMap()

[imaihal]

mlir/lib/Transforms/Utils/Utils.cpp

395

@bondhugula Thanks for review again. I updated the condition for tiled-layout map. How do you think of it?

[imaihal]

Updated to solve warning message in clang-tidy and rebased.

[imaihal]

Rebased.

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

406

I think this variable is misnamed - it should be modExprs?

407

Nit: you can use llvm::reverse(map.getResults()).

423–427

Shouldn't all the floordiv exprs come before mod's?

433

Range-based loop?

434–436

This can be written in one line using llvm::is_contained(floorDivs, ...).

[imaihal]

Reflected comments for isTiledLayoutMap().

[imaihal]

mlir/lib/Transforms/Utils/Utils.cpp

433

All of elements in modExprs need to be found in floorDivs. So, I think we need to check all of the elements.

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

432

Nit: avoid auto.

433

Yes, I was only asking for a switch to a range-based for loop from an iterator one. You could check all elements.

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

396

I'm sorry but what you have here again is a necessary condition but not sufficient. For eg. the following is not the kind of tiled layout maps you want to detect:
d0 / 32, d1 / 32, d0 / 64, d1 mod 32)
The sufficient condition is actually much simpler - you need a specific k dimensions being floordiv'ed and then those same k dimensions appearing with a mod w.r.t the same tile sizes and nothing else (no other expressions) involving those dimensions. Also, this function should provide those tile sizes back to the call site (as an extra output argument) - since a user would in nearly all cases want those for further action.

517–519

But looks like you are assuming here that every floordiv has a corresponding mod which isn't the case given the check you have. Consider
d0 / 32, d1 / 32, d0 / 64, d0 mod 32, d1 mod 32 - will this work correctly?
I think isTiledLayoutMap should just return the tile sizes and the corresponding positions (since the tile-space ordering can be different from the intra-tile one), and you could use it here. This way you would unify the logic for createDimSizeExprForTiledLayout with isTiledLayoutMap?

641

Nit: is -> is a

[imaihal]

mlir/lib/Transforms/Utils/Utils.cpp

396

@bondhugula Thanks for your review. Sorry, I may be confused about tiled layout.

• d0 / 32, d1 / 32, d0 / 64, d1 mod 32)

I implemented the function returns true for this map because d1 dimension meets tiled layout pattern (tiled in d1 dimension using tile size of 32, and not tiled in d0 dimension), but you suggested this is not true.
In this map, d0 dimension is floordived, but no mod for d0. Is this the reason why we can not say this is the tiled layout?

• d0 / 32, d1 / 32, d0 / 64, d0 mod 32, d1 mod 32

In another comment, you showed this map. Can we say this is the tiled layout ? (tiled in d0 dimension using tile size of 32, and tiled in d1 dimension using tile size of 32?)

[imaihal]

mlir/lib/Transforms/Utils/Utils.cpp

517–519

@bondhugula

d0 / 32, d1 / 32, d0 / 64, d0 mod 32, d1 mod 32 - will this work correctly?

This does not work correctly in current implementation. I will update using the tile sizes and the positions.

[imaihal]

Updated check routine for a tiled layout. It checks dimensions being floordiv'ed and then those same dimensions appearing with a mod and no other expressions involving those dimensions. Also, the routine provides those tile sizes back to the call site.

[imaihal]

Fixed errors in clang-tidy.

[imaihal]

@bondhugula I updated the patch.

The sufficient condition is actually much simpler - you need a specific k dimensions being floordiv'ed and then those same k dimensions appearing with a mod w.r.t the same tile sizes and nothing else (no other expressions) involving those dimensions. Also, this function should provide those tile sizes back to the call site (as an extra output argument) - since a user would in nearly all cases want those for further action.

I implemented this and changed the function name from isTiledLayoutMap() to getTileSizePos(). I hope this meets your request.

[imaihal]

Rebased.

[imaihal]

Hi, @bondhugula , I updated this patch to reflect your suggestions. How do you think of this?

[bondhugula]

In D97655#2736378, @imaihal wrote:

Hi, @bondhugula , I updated this patch to reflect your suggestions. How do you think of this?

Sorry about the delay. I should be able to review this (will hopefully be the final round) in the next one day.

mlir/lib/Transforms/Utils/Utils.cpp

399–400

Can you use an output parameter instead of returning? Sometimes RVO may not happen and will lead to an extra copy.

496

Nit: Can you add the output parameter as the last argument?

[imaihal]

Used output parameters to reflect comments.

[imaihal]

Rebased and removed unnecessary include file.

[bondhugula]

mlir/lib/Transforms/NormalizeMemRefs.cpp

42 ↗

(On Diff #342725)

This can be instead specified in a single line in Passes.td.

mlir/lib/Transforms/Utils/Utils.cpp

395

Many typos here:
no other expression doesn't involve -> no other expression involves
are appeared -> appear
by a tile size <-- this is inaccurate since tile sizes can be different

-> by respective tile sizes

398

an empty vector is returned.

422

It's more readable to return an empty vector - does {} or SmallVector<...>{}/() work?

478

Why do you need a separate vector just to return empty? Does {} or SmallVector<...>{}/() work?

513

This isn't clear - which dimension size is being referred to here?

641

single -> single level ?

703

tileSizePos.empty() instead of size() == 0.

713–719

if shape[d] < 0, you are adding an invalid constraint here: 0 <= s <= -2

[imaihal]

Reflected comments.

[imaihal]

@bondhugula Thanks for the review! I updated the code. Do you have any other comments?

mlir/lib/Transforms/Utils/Utils.cpp

513

This function checks if dim of normalized memref is dynamic using inMemrefTypeDynDims. Return true if dim is dynamic.

713–719

Changed to use the constraint only in static dimension.

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

517–519

You probably don't need to handle this unless there is a use case. So long as the form handled is documented, which I think you have, it should be fine.

[bondhugula]

LGTM

[imaihal]

@bondhugula Much appreciate your reviews!

I don't have commit access. Could you commit this patch?

[imaihal]

Rebased.
Could you land this patch?

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_CastOp : MemRef_Op<"cast", [
       NoSideEffect, SameOperandsAndResultShape,
       DeclareOpInterfaceMethods<CastOpInterface>,
-      ViewLikeOpInterface
+      ViewLikeOpInterface,
+      MemRefsNormalizable
     ]> {
   let summary = "memref cast operation";
   let description = [{
     Syntax:
 
     ```
     operation ::= ssa-id `=` `memref.cast` ssa-use `:` type `to` type
     ```
def MemRef_DeallocOp : MemRef_Op<"dealloc", [MemRefsNormalizable]> {
   let verifier = ?;
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // DimOp
 //===----------------------------------------------------------------------===//
 
-def DimOp : MemRef_Op<"dim", [NoSideEffect]> {
+def DimOp : MemRef_Op<"dim", [NoSideEffect, MemRefsNormalizable]> {
   let summary = "dimension index operation";
   let description = [{
     The `dim` operation takes a memref and a dimension operand of type `index`.
     It returns the size of the requested dimension of the given memref.
     If the dimension index is out of bounds the behavior is undefined.
 
     The specified memref type is that of the first operand.
 

mlir/include/mlir/Transforms/Passes.td

affine.for %arg3 = %c0 to %c8 {
         }
       }
     }
     return
   }
   ```
   }];
   let constructor = "mlir::createNormalizeMemRefsPass()";
+  let dependentDialects = ["AffineDialect"];
 }
 
 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",

mlir/lib/Transforms/PassDetail.h

 //===- PassDetail.h - Transforms Pass class details -------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TRANSFORMS_PASSDETAIL_H_
 #define TRANSFORMS_PASSDETAIL_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+class AffineDialect;
+
 // Forward declaration from Dialect.h
 template <typename ConcreteDialect>
 void registerDialect(DialectRegistry &registry);
 
 namespace linalg {
 class LinalgDialect;
 } // end namespace linalg
 

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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements miscellaneous transformation routines for non-loop IR
 // structures.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/Utils.h"
-
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dominance.h"
if (j < subOperands.size()) {
       newOperands[i] = (*sliceOps)[j];
     }
   }
   for (unsigned idx = 0, e = newOperands.size(); idx < e; idx++) {
     opInst->setOperand(idx, newOperands[idx]);
   }
 }
 
+/// Enum to set patterns of affine expr in tiled-layout map.
+/// TileFloorDiv: <dim expr> div <tile size>
+/// TileMod: <dim expr> mod <tile size>
+/// TileNone: None of the above
+/// Example:
+/// #tiled_2d_128x256 = affine_map<(d0, d1)
+///            -> (d0 div 128, d1 div 256, d0 mod 128, d1 mod 256)>
+/// "d0 div 128" and "d1 div 256" ==> TileFloorDiv
+/// "d0 mod 128" and "d1 mod 256" ==> TileMod
+enum TileExprPattern { TileFloorDiv, TileMod, TileNone };
+
+/// Check if `map` is a tiled layout. In the tiled layout, specific k dimensions
+/// being floordiv'ed by respective tile sizes appeare in a mod with the same
+/// tile sizes, and no other expression involves those k dimensions. This
+/// function stores a vector of tuples (`tileSizePos`) including AffineExpr for
+/// tile size, positions of corresponding `floordiv` and `mod`. If it is not a
+/// tiled layout, an empty vector is returned.
+static LogicalResult getTileSizePos(
+    AffineMap map,
+    SmallVectorImpl<std::tuple<AffineExpr, unsigned, unsigned>> &tileSizePos) {
+  // Create `floordivExprs` which is a vector of tuples including LHS and RHS of
+  // `floordiv` and its position in `map` output.
+  // Example: #tiled_2d_128x256 = affine_map<(d0, d1)
+  //                -> (d0 div 128, d1 div 256, d0 mod 128, d1 mod 256)>
+  // In this example, `floordivExprs` includes {d0, 128, 0} and {d1, 256, 1}.
+  SmallVector<std::tuple<AffineExpr, AffineExpr, unsigned>, 4> floordivExprs;
+  unsigned pos = 0;
+  for (AffineExpr expr : map.getResults()) {
+    if (expr.getKind() == AffineExprKind::FloorDiv) {
+      AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+      if (binaryExpr.getRHS().isa<AffineConstantExpr>())
+        floordivExprs.emplace_back(
+            std::make_tuple(binaryExpr.getLHS(), binaryExpr.getRHS(), pos));
+    }
+    pos++;
+  }
+  // Not tiled layout if `floordivExprs` is empty.
+  if (floordivExprs.empty()) {
+    tileSizePos = SmallVector<std::tuple<AffineExpr, unsigned, unsigned>>{};
+    return success();
+  }
+
+  // Check if LHS of `floordiv` is used in LHS of `mod`. If not used, `map` is
+  // not tiled layout.
+  for (std::tuple<AffineExpr, AffineExpr, unsigned> fexpr : floordivExprs) {
+    AffineExpr floordivExprLHS = std::get<0>(fexpr);
+    AffineExpr floordivExprRHS = std::get<1>(fexpr);
+    unsigned floordivPos = std::get<2>(fexpr);
+
+    // Walk affinexpr of `map` output except `fexpr`, and check if LHS and RHS
+    // of `fexpr` are used in LHS and RHS of `mod`. If LHS of `fexpr` is used
+    // other expr, the map is not tiled layout. Example of non tiled layout:
+    //   affine_map<(d0, d1, d2) -> (d0, d1, d2 floordiv 256, d2 floordiv 256)>
+    //   affine_map<(d0, d1, d2) -> (d0, d1, d2 floordiv 256, d2 mod 128)>
+    //   affine_map<(d0, d1, d2) -> (d0, d1, d2 floordiv 256, d2 mod 256, d2 mod
+    //   256)>
+    bool found = false;
+    pos = 0;
+    for (AffineExpr expr : map.getResults()) {
+      bool notTiled = false;
+      if (pos != floordivPos) {
+        expr.walk([&](AffineExpr e) {
+          if (e == floordivExprLHS) {
+            if (expr.getKind() == AffineExprKind::Mod) {
+              AffineBinaryOpExpr binaryExpr = expr.cast<AffineBinaryOpExpr>();
+              // If LHS and RHS of `mod` are the same with those of floordiv.
+              if (floordivExprLHS == binaryExpr.getLHS() &&
+                  floordivExprRHS == binaryExpr.getRHS()) {
+                // Save tile size (RHS of `mod`), and position of `floordiv` and
+                // `mod` if same expr with `mod` is not found yet.
+                if (!found) {
+                  tileSizePos.emplace_back(
+                      std::make_tuple(binaryExpr.getRHS(), floordivPos, pos));
+                  found = true;
+                } else {
+                  // Non tiled layout: Have multilpe `mod` with the same LHS.
+                  // eg. affine_map<(d0, d1, d2) -> (d0, d1, d2 floordiv 256, d2
+                  // mod 256, d2 mod 256)>
+                  notTiled = true;
+                }
+              } else {
+                // Non tiled layout: RHS of `mod` is different from `floordiv`.
+                // eg. affine_map<(d0, d1, d2) -> (d0, d1, d2 floordiv 256, d2
+                // mod 128)>
+                notTiled = true;
+              }
+            } else {
+              // Non tiled layout: LHS is the same, but not `mod`.
+              // eg. affine_map<(d0, d1, d2) -> (d0, d1, d2 floordiv 256, d2
+              // floordiv 256)>
+              notTiled = true;
+            }
+          }
+        });
+      }
+      if (notTiled) {
+        tileSizePos = SmallVector<std::tuple<AffineExpr, unsigned, unsigned>>{};
+        return success();
+      }
+      pos++;
+    }
+  }
+  return success();
+}
+
+/// Check if `dim` dimension of memrefType with `layoutMap` becomes dynamic
+/// after normalization. Dimensions that include dynamic dimensions in the map
+/// output will become dynamic dimensions. Return true if `dim` is dynamic
+/// dimension.
+///
+/// Example:
+/// #map0 = affine_map<(d0, d1) -> (d0, d1 floordiv 32, d1 mod 32)>
+///
+/// If d1 is dynamic dimension, 2nd and 3rd dimension of map output are dynamic.
+/// memref<4x?xf32, #map0>  ==>  memref<4x?x?xf32>
+static bool
+isNormalizedMemRefDynamicDim(unsigned dim, AffineMap layoutMap,
+                             SmallVectorImpl<unsigned> &inMemrefTypeDynDims,
+                             MLIRContext *context) {
+  bool isDynamicDim = false;
+  AffineExpr expr = layoutMap.getResults()[dim];
+  // Check if affine expr of the dimension includes dynamic dimension of input
+  // memrefType.
+  expr.walk([&inMemrefTypeDynDims, &isDynamicDim, &context](AffineExpr e) {
+    if (e.isa<AffineDimExpr>()) {
+      for (unsigned dm : inMemrefTypeDynDims) {
+        if (e == getAffineDimExpr(dm, context)) {
+          isDynamicDim = true;
+        }
+      }
+    }
+  });
+  return isDynamicDim;
+}
+
+/// Create affine expr to calculate dimension size for a tiled-layout map.
+static AffineExpr createDimSizeExprForTiledLayout(AffineExpr oldMapOutput,
+                                                  TileExprPattern pat) {
+  // Create map output for the patterns.
+  // "floordiv <tile size>" ==> "ceildiv <tile size>"
+  // "mod <tile size>" ==> "<tile size>"
+  AffineExpr newMapOutput;
+  AffineBinaryOpExpr binaryExpr = nullptr;
+  switch (pat) {
+  case TileExprPattern::TileMod:
+    binaryExpr = oldMapOutput.cast<AffineBinaryOpExpr>();
+    newMapOutput = binaryExpr.getRHS();
+    break;
+  case TileExprPattern::TileFloorDiv:
+    binaryExpr = oldMapOutput.cast<AffineBinaryOpExpr>();
+    newMapOutput = getAffineBinaryOpExpr(
+        AffineExprKind::CeilDiv, binaryExpr.getLHS(), binaryExpr.getRHS());
+    break;
+  default:
+    newMapOutput = oldMapOutput;
+  }
+  return newMapOutput;
+}
+
+/// Create new maps to calculate each dimension size of `newMemRefType`, and
+/// create `newDynamicSizes` from them by using AffineApplyOp.
+///
+/// Steps for normalizing dynamic memrefs for a tiled layout map
+/// Example:
+///    #map0 = affine_map<(d0, d1) -> (d0, d1 floordiv 32, d1 mod 32)>
+///    %0 = dim %arg0, %c1 :memref<4x?xf32>
+///    %1 = alloc(%0) : memref<4x?xf32, #map0>
+///
+/// (Before this function)
+/// 1. Check if `map`(#map0) is a tiled layout using `getTileSizePos()`. Only
+/// single layout map is supported.
+///
+/// 2. Create normalized memrefType using `isNormalizedMemRefDynamicDim()`. It
+/// is memref<4x?x?xf32> in the above example.
+///
+/// (In this function)
+/// 3. Create new maps to calculate each dimension of the normalized memrefType
+/// using `createDimSizeExprForTiledLayout()`. In the tiled layout, the
+/// dimension size can be calculated by replacing "floordiv <tile size>" with
+/// "ceildiv <tile size>" and "mod <tile size>" with "<tile size>".
+/// - New map in the above example
+///   #map0 = affine_map<(d0, d1) -> (d0)>
+///   #map1 = affine_map<(d0, d1) -> (d1 ceildiv 32)>
+///   #map2 = affine_map<(d0, d1) -> (32)>
+///
+/// 4. Create AffineApplyOp to apply the new maps. The output of AffineApplyOp
+/// is used in dynamicSizes of new AllocOp.
+///   %0 = dim %arg0, %c1 : memref<4x?xf32>
+///   %c4 = constant 4 : index
+///   %1 = affine.apply #map1(%c4, %0)
+///   %2 = affine.apply #map2(%c4, %0)
+static void createNewDynamicSizes(MemRefType oldMemRefType,
+                                  MemRefType newMemRefType, AffineMap map,
+                                  memref::AllocOp *allocOp, OpBuilder b,
+                                  SmallVectorImpl<Value> &newDynamicSizes) {
+  // Create new input for AffineApplyOp.
+  SmallVector<Value, 4> inAffineApply;
+  ArrayRef<int64_t> oldMemRefShape = oldMemRefType.getShape();
+  unsigned dynIdx = 0;
+  for (unsigned d = 0; d < oldMemRefType.getRank(); ++d) {
+    if (oldMemRefShape[d] < 0) {
+      // Use dynamicSizes of allocOp for dynamic dimension.
+      inAffineApply.emplace_back(allocOp->dynamicSizes()[dynIdx]);
+      dynIdx++;
+    } else {
+      // Create ConstantOp for static dimension.
+      Attribute constantAttr =
+          b.getIntegerAttr(b.getIndexType(), oldMemRefShape[d]);
+      inAffineApply.emplace_back(
+          b.create<ConstantOp>(allocOp->getLoc(), constantAttr));
+    }
+  }
+
+  // Create new map to calculate each dimension size of new memref for each
+  // original map output. Only for dynamic dimesion of `newMemRefType`.
+  unsigned newDimIdx = 0;
+  ArrayRef<int64_t> newMemRefShape = newMemRefType.getShape();
+  SmallVector<std::tuple<AffineExpr, unsigned, unsigned>> tileSizePos;
+  (void)getTileSizePos(map, tileSizePos);
+  for (AffineExpr expr : map.getResults()) {
+    if (newMemRefShape[newDimIdx] < 0) {
+      // Create new maps to calculate each dimension size of new memref.
+      enum TileExprPattern pat = TileExprPattern::TileNone;
+      for (auto pos : tileSizePos) {
+        if (newDimIdx == std::get<1>(pos))
+          pat = TileExprPattern::TileFloorDiv;
+        else if (newDimIdx == std::get<2>(pos))
+          pat = TileExprPattern::TileMod;
+      }
+      AffineExpr newMapOutput = createDimSizeExprForTiledLayout(expr, pat);
+      AffineMap newMap =
+          AffineMap::get(map.getNumInputs(), map.getNumSymbols(), newMapOutput);
+      Value affineApp =
+          b.create<AffineApplyOp>(allocOp->getLoc(), newMap, inAffineApply);
+      newDynamicSizes.emplace_back(affineApp);
+    }
+    newDimIdx++;
+  }
+}
+
 // TODO: Currently works for static memrefs with a single layout map.
 LogicalResult mlir::normalizeMemRef(memref::AllocOp *allocOp) {
   MemRefType memrefType = allocOp->getType();
   OpBuilder b(*allocOp);
 
   // Fetch a new memref type after normalizing the old memref to have an
   // identity map layout.
   MemRefType newMemRefType =
       normalizeMemRefType(memrefType, b, allocOp->symbolOperands().size());
   if (newMemRefType == memrefType)
     // Either memrefType already had an identity map or the map couldn't be
     // transformed to an identity map.
     return failure();
 
   Value oldMemRef = allocOp->getResult();
 
   SmallVector<Value, 4> symbolOperands(allocOp->symbolOperands());
-  memref::AllocOp newAlloc = b.create<memref::AllocOp>(
-      allocOp->getLoc(), newMemRefType, allocOp->alignmentAttr());
   AffineMap layoutMap = memrefType.getAffineMaps().front();
+  memref::AllocOp newAlloc;
+  // Check if `layoutMap` is a tiled layout. Only single layout map is
+  // supported for normalizing dynamic memrefs.
+  SmallVector<std::tuple<AffineExpr, unsigned, unsigned>> tileSizePos;
+  (void)getTileSizePos(layoutMap, tileSizePos);
+  if (newMemRefType.getNumDynamicDims() > 0 && !tileSizePos.empty()) {
+    MemRefType oldMemRefType = oldMemRef.getType().cast<MemRefType>();
+    SmallVector<Value, 4> newDynamicSizes;
+    createNewDynamicSizes(oldMemRefType, newMemRefType, layoutMap, allocOp, b,
+                          newDynamicSizes);
+    // Add the new dynamic sizes in new AllocOp.
+    newAlloc =
+        b.create<memref::AllocOp>(allocOp->getLoc(), newMemRefType,
+                                  newDynamicSizes, allocOp->alignmentAttr());
+  } else {
+    newAlloc = b.create<memref::AllocOp>(allocOp->getLoc(), newMemRefType,
+                                         allocOp->alignmentAttr());
+  }
   // Replace all uses of the old memref.
   if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
                                       /*extraIndices=*/{},
                                       /*indexRemap=*/layoutMap,
                                       /*extraOperands=*/{},
                                       /*symbolOperands=*/symbolOperands,
                                       /*domInstFilter=*/nullptr,
                                       /*postDomInstFilter=*/nullptr,
if (layoutMaps.empty() ||
     // 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)
+  // Normalize only static memrefs and dynamic memrefs with a tiled-layout map
+  // for now.
+  // TODO: Normalize the other types of dynamic memrefs.
+  SmallVector<std::tuple<AffineExpr, unsigned, unsigned>> tileSizePos;
+  (void)getTileSizePos(layoutMaps.front(), tileSizePos);
+  if (memrefType.getNumDynamicDims() > 0 && tileSizePos.empty())
     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.
   ArrayRef<int64_t> shape = memrefType.getShape();
   // FlatAffineConstraint may later on use symbolicOperands.
   FlatAffineConstraints fac(rank, numSymbolicOperands);
+  SmallVector<unsigned, 4> memrefTypeDynDims;
   for (unsigned d = 0; d < rank; ++d) {
+    // Use constraint system only in static dimensions.
+    if (shape[d] > 0) {
       fac.addConstantLowerBound(d, 0);
       fac.addConstantUpperBound(d, shape[d] - 1);
+    } else {
+      memrefTypeDynDims.emplace_back(d);
+    }
   }
   // 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)))
     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) {
+    // Check if each dimension of normalized memrefType is dynamic.
+    bool isDynDim = isNormalizedMemRefDynamicDim(
+        d, layoutMap, memrefTypeDynDims, b.getContext());
+    if (isDynDim) {
+      newShape[d] = -1;
+    } else {
       // 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.
       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 memrefType;
+      // If dimension of new memrefType is dynamic, the value is -1.
       newShape[d] = ubConst.getValue() + 1;
     }
+  }
 
   // Create the new memref type after trivializing the old layout map.
   MemRefType newMemRefType =
       MemRefType::Builder(memrefType)
           .setShape(newShape)
           .setAffineMaps(b.getMultiDimIdentityMap(newRank));
 
   return newMemRefType;
 }

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

+// RUN: mlir-opt -normalize-memrefs %s -split-input-file| FileCheck %s
+
+// For all these cases, we test if MemRefs Normalization works with the test
+// operations. These are test cases for MemRefs with dynamic dimension
+// and tiled-layout map.
+// * test.op_norm: this operation has the MemRefsNormalizable attribute. The tests
+//   that include this operation are constructed so that the normalization should
+//   happen.
+
+#map_tiled = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 64, d2 mod 32, d3 mod 64)>
+
+// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 ceildiv 32)>
+// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (32)>
+
+// Test with op_norm and maps in arguments and in the operations in the function.
+// Memref has two dynamic dimensions.
+
+// CHECK-LABEL:  test_norm_dynamic12
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x1x?x64xf32>) {
+func @test_norm_dynamic12(%arg0 : memref<1x?x?x14xf32, #map_tiled>) -> () {
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_tiled>
+    %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_tiled>
+    %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_tiled>
+    "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_tiled>, memref<1x?x?x14xf32, #map_tiled>) -> ()
+    memref.dealloc %2 :  memref<1x?x?x14xf32, #map_tiled>
+    return
+    // CHECK-DAG:       [[CST_1_:%.+]] = constant 1 : index
+    // CHECK-DAG:       [[CST_2_:%.+]] = constant 2 : index
+    // CHECK-NOT: separator of consecutive DAGs
+    // CHECK-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x1x?x64xf32>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x1x?x64xf32>
+    // CHECK-DAG:       [[CST_1_1_:%.+]] = constant 1 : index
+    // CHECK-DAG:       [[CST_14_:%.+]] = constant 14 : index
+    // CHECK-NOT: separator of consecutive DAGs
+    // CHECK-DAG:       [[VAR_2_:%.+]] = affine.apply #[[$MAP0]]([[CST_1_1_]], [[DIM_0_]], [[DIM_1_]], [[CST_14_]])
+    // CHECK-DAG:       [[VAR_3_:%.+]] = affine.apply #[[$MAP1]]([[CST_1_1_]], [[DIM_0_]], [[DIM_1_]], [[CST_14_]])
+    // CHECK-DAG:       [[VAR_4_:%.+]] = affine.apply #[[$MAP2]]([[CST_1_1_]], [[DIM_0_]], [[DIM_1_]], [[CST_14_]])
+    // CHECK:           [[RES_:%.+]] = memref.alloc([[VAR_2_]], [[VAR_3_]], [[VAR_4_]]) : memref<1x?x?x1x?x64xf32>
+    // CHECK:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x1x?x64xf32>, memref<1x?x?x1x?x64xf32>) -> ()
+    // CHECK:           memref.dealloc [[RES_]] : memref<1x?x?x1x?x64xf32>
+    // CHECK:           return
+}
+
+// -----
+
+// Test with op_norm and maps in arguments and in the operations in the function.
+// All of dimensions are dynamic.
+
+#map_tiled1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, (d2 floordiv 4) floordiv 32, (d3 mod 8) floordiv 64, (d2 floordiv 4) mod 32, (d3 mod 8) mod 64)>
+
+// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> ((d2 floordiv 4) ceildiv 32)>
+// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (32)>
+// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>
+// CHECK-DAG: #[[$MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> ((d3 mod 8) ceildiv 64)>
+// CHECK-DAG: #[[$MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (64)>
+
+// CHECK-LABEL:  test_norm_dynamic1234
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<?x?x?x?x?x?xf32>) {
+func @test_norm_dynamic1234(%arg0 : memref<?x?x?x?xf32, #map_tiled1>) -> () {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %c3 = constant 3 : index
+    %0 = memref.dim %arg0, %c0 :memref<?x?x?x?xf32, #map_tiled1>
+    %1 = memref.dim %arg0, %c1 :memref<?x?x?x?xf32, #map_tiled1>
+    %2 = memref.dim %arg0, %c2 :memref<?x?x?x?xf32, #map_tiled1>
+    %3 = memref.dim %arg0, %c3 :memref<?x?x?x?xf32, #map_tiled1>
+    %4 = memref.alloc(%0, %1, %2, %3) : memref<?x?x?x?xf32, #map_tiled1>
+    "test.op_norm"(%arg0, %4) : (memref<?x?x?x?xf32, #map_tiled1>, memref<?x?x?x?xf32, #map_tiled1>) -> ()
+    memref.dealloc %4 :  memref<?x?x?x?xf32, #map_tiled1>
+    return
+    // CHECK-DAG:       [[CST_0_:%.+]] = constant 0 : index
+    // CHECK-DAG:       [[CST_1_:%.+]] = constant 1 : index
+    // CHECK-DAG:       [[CST_2_:%.+]] = constant 2 : index
+    // CHECK-DAG:       [[CST_3_:%.+]] = constant 3 : index
+    // CHECK-NOT: separator of consecutive DAGs
+    // CHECK-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_0_]] : memref<?x?x?x?x?x?xf32>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<?x?x?x?x?x?xf32>
+    // CHECK-DAG:       [[DIM_2_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<?x?x?x?x?x?xf32>
+    // CHECK-DAG:       [[DIM_3_:%.+]] = memref.dim [[ARG_0_]], [[CST_3_]] : memref<?x?x?x?x?x?xf32>
+    // CHECK-NOT: separator of consecutive DAGs
+    // CHECK-DAG:       [[VAR_4_:%.+]] = affine.apply #[[$MAP3]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])
+    // CHECK-DAG:       [[VAR_5_:%.+]] = affine.apply #[[$MAP0]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])
+    // CHECK-DAG:       [[VAR_6_:%.+]] = affine.apply #[[$MAP1]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])
+    // CHECK-DAG:       [[VAR_7_:%.+]] = affine.apply #[[$MAP4]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])
+    // CHECK-DAG:       [[VAR_8_:%.+]] = affine.apply #[[$MAP2]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])
+    // CHECK-DAG:       [[VAR_9_:%.+]] = affine.apply #[[$MAP5]]([[DIM_0_]], [[DIM_1_]], [[DIM_2_]], [[DIM_3_]])
+    // CHECK:           [[RES_:%.+]] = memref.alloc([[VAR_4_]], [[VAR_5_]], [[VAR_6_]], [[VAR_7_]], [[VAR_8_]], [[VAR_9_]]) : memref<?x?x?x?x?x?xf32>
+    // CHECK:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<?x?x?x?x?x?xf32>, memref<?x?x?x?x?x?xf32>) -> ()
+    // CHECK:           memref.dealloc [[RES_]] : memref<?x?x?x?x?x?xf32>
+    // CHECK:           return
+}
+
+// -----
+
+// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.
+// This is not normalized since this is not tiled-layout map. No mod and floordiv.
+
+#map_not_tiled0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2)>
+
+// CHECK-DAG: #[[$MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2)>
+
+// CHECK-LABEL:  func @test_norm_dynamic_not_tiled0
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP6]]>) {
+func @test_norm_dynamic_not_tiled0(%arg0 : memref<1x?x?x14xf32, #map_not_tiled0>) -> () {
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled0>
+    %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled0>
+    %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled0>
+    "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled0>, memref<1x?x?x14xf32, #map_not_tiled0>) -> ()
+    memref.dealloc %2 :  memref<1x?x?x14xf32, #map_not_tiled0>
+    return
+    // CHECK-DAG:       [[CST_1_:%.+]] = constant 1 : index
+    // CHECK-DAG:       [[CST_2_:%.+]] = constant 2 : index
+    // CHECK-NOT: separator of consecutive DAGs
+    // CHECK-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK:           [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP6]]>, memref<1x?x?x14xf32, #[[$MAP6]]>) -> ()
+    // CHECK:           memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK:           return
+}
+
+// -----
+
+// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.
+// This is not normalized since this is not tiled-layout map. No floordiv.
+
+#map_not_tiled1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2, d2 mod 32, d3 mod 64)>
+
+// CHECK-DAG: #[[$MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2, d2 mod 32, d3 mod 64)>
+
+// CHECK-LABEL:  func @test_norm_dynamic_not_tiled1
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP6]]>) {
+func @test_norm_dynamic_not_tiled1(%arg0 : memref<1x?x?x14xf32, #map_not_tiled1>) -> () {
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled1>
+    %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled1>
+    %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled1>
+    "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled1>, memref<1x?x?x14xf32, #map_not_tiled1>) -> ()
+    memref.dealloc %2 :  memref<1x?x?x14xf32, #map_not_tiled1>
+    return
+    // CHECK-DAG:       [[CST_1_:%.+]] = constant 1 : index
+    // CHECK-DAG:       [[CST_2_:%.+]] = constant 2 : index
+    // CHECK-NOT: separator of consecutive DAGs
+    // CHECK-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK:           [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP6]]>, memref<1x?x?x14xf32, #[[$MAP6]]>) -> ()
+    // CHECK:           memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK:           return
+}
+
+// -----
+
+// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.
+// This is not normalized since this is not tiled-layout map. RHS of floordiv is different from that of mod.
+
+#map_not_tiled2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 floordiv 64, d2 mod 32, d3 mod 32)>
+
+// CHECK-DAG: #[[$MAP7:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 floordiv 64, d2 mod 32, d3 mod 32)>
+
+// CHECK-LABEL:  func @test_norm_dynamic_not_tiled2
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP7]]>) {
+func @test_norm_dynamic_not_tiled2(%arg0 : memref<1x?x?x14xf32, #map_not_tiled2>) -> () {
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled2>
+    %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled2>
+    %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled2>
+    "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled2>, memref<1x?x?x14xf32, #map_not_tiled2>) -> ()
+    memref.dealloc %2 :  memref<1x?x?x14xf32, #map_not_tiled2>
+    return
+    // CHECK-DAG:       [[CST_1_:%.+]] = constant 1 : index
+    // CHECK-DAG:       [[CST_2_:%.+]] = constant 2 : index
+    // CHECK-NOT: separator of consecutive DAGs
+    // CHECK-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP7]]>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP7]]>
+    // CHECK:           [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP7]]>
+    // CHECK:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP7]]>, memref<1x?x?x14xf32, #[[$MAP7]]>) -> ()
+    // CHECK:           memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP7]]>
+    // CHECK:           return
+}
+
+// -----
+
+// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.
+// This is not normalized since this is not tiled-layout map. Multiple mod with the same LHS and RHS.
+
+#map_not_tiled3 = affine_map<(d0, d1, d2, d3) -> (d0, d1 floordiv 32, d2, d3, d1 mod 32, d1 mod 32)>
+
+// CHECK-DAG: #[[$MAP8:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 floordiv 32, d2, d3, d1 mod 32, d1 mod 32)>
+
+// CHECK-LABEL:  func @test_norm_dynamic_not_tiled3
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP8]]>) {
+func @test_norm_dynamic_not_tiled3(%arg0 : memref<1x?x?x14xf32, #map_not_tiled3>) -> () {
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled3>
+    %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled3>
+    %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled3>
+    "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled3>, memref<1x?x?x14xf32, #map_not_tiled3>) -> ()
+    memref.dealloc %2 :  memref<1x?x?x14xf32, #map_not_tiled3>
+    return
+    // CHECK-DAG:       [[CST_1_:%.+]] = constant 1 : index
+    // CHECK-DAG:       [[CST_2_:%.+]] = constant 2 : index
+    // CHECK-NOT: separator of consecutive DAGs
+    // CHECK-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP8]]>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP8]]>
+    // CHECK:           [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP8]]>
+    // CHECK:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP8]]>, memref<1x?x?x14xf32, #[[$MAP8]]>) -> ()
+    // CHECK:           memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP8]]>
+    // CHECK:           return
+}
+
+// -----
+
+// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.
+// This is not normalized since this is not tiled-layout map. floordiv and mod with the same LHS and RHS(d0 floordiv 32 and d0 mod 32), but, unrelaed d0 exists in other position.
+
+#map_not_tiled4 = affine_map<(d0, d1, d2, d3) -> (d0 floordiv 32, d1 floordiv 32, d0, d3, d0 mod 32, d1 mod 32)>
+
+// CHECK-DAG: #[[$MAP9:.+]] = affine_map<(d0, d1, d2, d3) -> (d0 floordiv 32, d1 floordiv 32, d0, d3, d0 mod 32, d1 mod 32)>
+
+// CHECK-LABEL:  func @test_norm_dynamic_not_tiled4
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP9]]>) {
+func @test_norm_dynamic_not_tiled4(%arg0 : memref<1x?x?x14xf32, #map_not_tiled4>) -> () {
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %0 = memref.dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled4>
+    %1 = memref.dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled4>
+    %2 = memref.alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled4>
+    "test.op_norm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled4>, memref<1x?x?x14xf32, #map_not_tiled4>) -> ()
+    memref.dealloc %2 :  memref<1x?x?x14xf32, #map_not_tiled4>
+    return
+    // CHECK-DAG:       [[CST_1_:%.+]] = constant 1 : index
+    // CHECK-DAG:       [[CST_2_:%.+]] = constant 2 : index
+    // CHECK-NOT: separator of consecutive DAGs
+    // CHECK-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP9]]>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP9]]>
+    // CHECK:           [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP9]]>
+    // CHECK:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP9]]>, memref<1x?x?x14xf32, #[[$MAP9]]>) -> ()
+    // CHECK:           memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP9]]>
+    // CHECK:           return
+}