[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.

598–600

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

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

242 ↗

(On Diff #335217)

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

391 ↗

(On Diff #335217)

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?

591

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

41

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?

591

single -> single level ?

649

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

659–661

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.

659–661

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/StandardOps/IR/Ops.td

def DeallocOp : Std_Op<"dealloc", [MemRefsNormalizable]> {
   let hasFolder = 1;
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // DimOp
 //===----------------------------------------------------------------------===//
 
-def DimOp : Std_Op<"dim", [NoSideEffect]> {
+def DimOp : Std_Op<"dim", [NoSideEffect, MemRefsNormalizable]> {
   let summary = "dimension index operation";
   let description = [{
     The `dim` operation takes a memref/tensor and a dimension operand of type
     `index`.
     It returns the size of the requested dimension of the given memref/tensor.
     If the dimension index is out of bounds the behavior is undefined.
 
     The specified memref or tensor type is that of the first operand.

mlir/lib/Transforms/NormalizeMemRefs.cpp

 /// non-normalizable as well. We assume external functions to be normalizable.
 struct NormalizeMemRefs : public NormalizeMemRefsBase<NormalizeMemRefs> {
   void runOnOperation() override;
   void normalizeFuncOpMemRefs(FuncOp funcOp, ModuleOp moduleOp);
   bool areMemRefsNormalizable(FuncOp funcOp);
   void updateFunctionSignature(FuncOp funcOp, ModuleOp moduleOp);
   void setCalleesAndCallersNonNormalizable(FuncOp funcOp, ModuleOp moduleOp,
                                            DenseSet<FuncOp> &normalizableFuncs);
+  // The affine.apply op is needed to normalize dynamic memrefs.
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect>();
+  }
   Operation *createOpResultsNormalized(FuncOp funcOp, Operation *oldOp);
 };
 
 } // end anonymous namespace
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createNormalizeMemRefsPass() {
   return std::make_unique<NormalizeMemRefs>();
 }

mlir/lib/Transforms/Utils/Utils.cpp

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.
+static bool isTiledLayoutMap(AffineMap map) {
+  // Find `floordiv` in `map` and get tile size.
+  // Example:
+  // #tiled_2d_128x256 = affine_map<(d0, d1)
+  //            -> (d0 div 128, d1 div 256, d0 mod 128, d1 mod 256)>
+  // In this example, `tileSizes` includes {d0, 128} and {d1, 256}.
+  SmallVector<std::pair<AffineExpr, AffineExpr>, 4> tileSizes;
+  for (AffineExpr expr : map.getResults()) {
+    expr.walk([&](AffineExpr e) {
+      if (e.getKind() == AffineExprKind::FloorDiv) {
+        AffineBinaryOpExpr binaryExpr = e.cast<AffineBinaryOpExpr>();
+        if (binaryExpr.getRHS().isa<AffineConstantExpr>())
+          tileSizes.emplace_back(binaryExpr.getLHS(), binaryExpr.getRHS());
+      }
+    });
+  }
+  if (tileSizes.empty())
+    return false;
+
+  // Find `mod` in `map` and check if it matches tiled layout pattern.
+  unsigned tileCount = 0;
+  for (AffineExpr expr : map.getResults()) {
+    expr.walk([&](AffineExpr e) {
+      if (e.getKind() == AffineExprKind::Mod) {
+        AffineBinaryOpExpr binaryExpr = e.cast<AffineBinaryOpExpr>();
+        for (const std::pair<AffineExpr, AffineExpr> &ts : tileSizes) {
+          if (binaryExpr.getLHS() == ts.first &&
+              binaryExpr.getRHS() == ts.second)
+            tileCount++;
+        }
+      }
+    });
+  }
+  if (tileCount == tileSizes.size())
+    return true;
+  return false;
+}
+
+/// Check if `dim` dimension of memrefType with `layoutMap` after normalization
+/// is dynamic. Dimensions that include dynamic dimensions in the map output
+/// will be dynamic dimensions.
+///
+/// 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 newDynDim = false;
+  AffineExpr expr = layoutMap.getResults()[dim];
+  // Check if affine expr of the dimension includes dynamic dimension of input
+  // memrefType.
+  expr.walk([&inMemrefTypeDynDims, &newDynDim, &context](AffineExpr e) {
+    if (e.getKind() == AffineExprKind::DimId) {
+      for (unsigned dm : inMemrefTypeDynDims) {
+        if (e == getAffineDimExpr(dm, context))
+          newDynDim = true;
+      }
+    }
+  });
+  return newDynDim;
+}
+
+/// Create affine expr to calculate dimension size for a tiled-layout map.
+static AffineExpr createDimSizeExprForTiledLayout(AffineExpr oldMapOutput) {
+  enum TileExprPattern pat = TileExprPattern::TileNone;
+  // Find "mod <tile size>" and "floordiv <tile size>" patterns.
+  oldMapOutput.walk([&pat](AffineExpr e) {
+    if (e.getKind() == AffineExprKind::Mod) {
+      AffineBinaryOpExpr binaryExpr = e.cast<AffineBinaryOpExpr>();
+      if (binaryExpr.getRHS().isa<AffineConstantExpr>())
+        pat = TileExprPattern::TileMod;
+    } else if (e.getKind() == AffineExprKind::FloorDiv) {
+      AffineBinaryOpExpr binaryExpr = e.cast<AffineBinaryOpExpr>();
+      if (binaryExpr.getRHS().isa<AffineConstantExpr>())
+        pat = TileExprPattern::TileFloorDiv;
+    }
+  });
+  // Create map output for the patterns.
+  // "floordiv <tile size>" ==> "ceildiv <tile size>"
+  // "mod <tile size>" ==> "<tile size>"
+  AffineExpr newMapOutput;
+  if (pat == TileExprPattern::TileMod) {
+    AffineBinaryOpExpr binaryExpr = oldMapOutput.cast<AffineBinaryOpExpr>();
+    newMapOutput = binaryExpr.getRHS();
+  } else if (pat == TileExprPattern::TileFloorDiv) {
+    AffineBinaryOpExpr binaryExpr = oldMapOutput.cast<AffineBinaryOpExpr>();
+    newMapOutput = getAffineBinaryOpExpr(
+        AffineExprKind::CeilDiv, binaryExpr.getLHS(), binaryExpr.getRHS());
+  } else {
+    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. Only single tiled layout is
+/// supported.(`isTiledLayoutMap()`))
+///
+/// 2. Create normalized memrefType(`newMemRefType`).
+/// (`isNormalizedMemRefDynamicDim()`)
+///
+/// (In this function)
+/// 3. Create new maps to calculate each dimension size of `newMemRefType`.
+/// Inputs of the new map are the same with `map`(d0, d1), and output of the new
+/// map is for each dimension. 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>".
+/// (`createDimSizeExprForTiledLayout()`)
+/// - New map
+///   #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 dynamicSizes for new AllocOp(`newDynamicSizes`). The inputs for
+/// AffineApplyOp are created from dynamicSizes of `allocOp` for dynamic
+/// dimension(%0) and created from ConstantOp for static dimension(%c4).
+///   %0 = dim %arg0, %c1 : memref<4x?xf32>
+///   %c4 = constant 4 : index
+///   %1 = affine.apply #map1(%c4, %0)
+///   %2 = affine.apply #map2(%c4, %0)
+///
+/// (After this function)
+/// 5. `newDynamicSizes`(%1 and %2) will be used in new AllocOp.
+/// (`normalizeMemRef()`)
+static void createNewDynamicSizes(MemRefType oldMemRefType,
+                                  MemRefType newMemRefType, AffineMap map,
+                                  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();
+  for (AffineExpr expr : map.getResults()) {
+    if (newMemRefShape[newDimIdx] < 0) {
+      // Create new maps to calculate each dimension size of new memref.
+      AffineExpr newMapOutput = createDimSizeExprForTiledLayout(expr);
+      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(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());
-  AllocOp newAlloc = b.create<AllocOp>(allocOp.getLoc(), newMemRefType,
-                                       allocOp.alignmentAttr());
   AffineMap layoutMap = memrefType.getAffineMaps().front();
+  AllocOp newAlloc;
+  // Check if `layoutMap` is tiled layout. Only single tiled layout is supported
+  // for normalizing dynamic memrefs.
+  if (newMemRefType.getNumDynamicDims() > 0 && isTiledLayoutMap(layoutMap)) {
+    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<AllocOp>(allocOp.getLoc(), newMemRefType,
+                                 newDynamicSizes, allocOp.alignmentAttr());
+  } else {
+    newAlloc = b.create<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.
+  if (memrefType.getNumDynamicDims() > 0 &&
+      !isTiledLayoutMap(layoutMaps.front()))
     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) {
     fac.addConstantLowerBound(d, 0);
     fac.addConstantUpperBound(d, shape[d] - 1);
+    if (shape[d] < 0)
+      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 newDynDim = isNormalizedMemRefDynamicDim(
+        d, layoutMap, memrefTypeDynDims, b.getContext());
     // 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)
+    if (ubConst.getValue() < 0 && memrefType.getNumDynamicDims() == 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.
+    if (newDynDim)
+      newShape[d] = -1;
+    else
       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 -canonicalize %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 = dim %arg0, %c1 :memref<1x?x?x14xf32, #map_tiled>
+    %1 = dim %arg0, %c2 :memref<1x?x?x14xf32, #map_tiled>
+    %2 = 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>) -> ()
+    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_:%.+]] = dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x1x?x64xf32>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = 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_:%.+]] = 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:           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_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)>
+// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)>
+// CHECK-DAG: #[[$MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d3 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_tiled>) -> () {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %c3 = constant 3 : index
+    %0 = dim %arg0, %c0 :memref<?x?x?x?xf32, #map_tiled>
+    %1 = dim %arg0, %c1 :memref<?x?x?x?xf32, #map_tiled>
+    %2 = dim %arg0, %c2 :memref<?x?x?x?xf32, #map_tiled>
+    %3 = dim %arg0, %c3 :memref<?x?x?x?xf32, #map_tiled>
+    %4 = alloc(%0, %1, %2, %3) : memref<?x?x?x?xf32, #map_tiled>
+    "test.op_norm"(%arg0, %4) : (memref<?x?x?x?xf32, #map_tiled>, memref<?x?x?x?xf32, #map_tiled>) -> ()
+    dealloc %4 :  memref<?x?x?x?xf32, #map_tiled>
+    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_:%.+]] = dim [[ARG_0_]], [[CST_0_]] : memref<?x?x?x?x?x?xf32>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = dim [[ARG_0_]], [[CST_1_]] : memref<?x?x?x?x?x?xf32>
+    // CHECK-DAG:       [[DIM_2_:%.+]] = dim [[ARG_0_]], [[CST_2_]] : memref<?x?x?x?x?x?xf32>
+    // CHECK-DAG:       [[DIM_3_:%.+]] = 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_:%.+]] = 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:           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 only tiled-layout map is supported.
+
+#map_not_tiled = 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_tiled
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP6]]>) {
+func @test_norm_dynamic_not_tiled(%arg0 : memref<1x?x?x14xf32, #map_not_tiled>) -> () {
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %0 = dim %arg0, %c1 :memref<1x?x?x14xf32, #map_not_tiled>
+    %1 = dim %arg0, %c2 :memref<1x?x?x14xf32, #map_not_tiled>
+    %2 = alloc(%0, %1) : memref<1x?x?x14xf32, #map_not_tiled>
+    "test.op_nonnorm"(%arg0, %2) : (memref<1x?x?x14xf32, #map_not_tiled>, memref<1x?x?x14xf32, #map_not_tiled>) -> ()
+    dealloc %2 :  memref<1x?x?x14xf32, #map_not_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_:%.+]] = dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK:           [[RES_:%.+]] = alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK:           "test.op_nonnorm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP6]]>, memref<1x?x?x14xf32, #[[$MAP6]]>) -> ()
+    // CHECK:           dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+    // CHECK:           return
+}