mlir: partially support normalization of memrefs with dynamic sizes

[rohany]

Summary

This patch adds support for partial normalization of dynamically sized
memrefs. In particular, it allows for normalization of memrefs that
appear as function arguments. It does not consider normalization of
memrefs allocated within functions, as that procedure requires
computation of the new dynamic sizes of the memrefs.

[rohany]

@bondhugula, a diff from the forum discussion we had earlier.

mlir/include/mlir/Dialect/Affine/Utils.h

 
 /// 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(memref::AllocOp *op);
 
 /// Normalizes `memrefType` so that the affine layout map of the memref is
 /// transformed to an identity map with a new shape being computed for the
-/// normalized memref type and returns it. The old memref type is simplify
+/// normalized memref type and returns it. The old memref type is simply
 /// returned if the normalization failed.
-MemRefType normalizeMemRefType(MemRefType memrefType,
+MemRefType normalizeMemRefType(OpBuilder& builder, MemRefType memrefType,
                                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

mlir/lib/Dialect/Affine/Utils/Utils.cpp

 // 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, allocOp->getSymbolOperands().size());
+      normalizeMemRefType(b, memrefType, allocOp->getSymbolOperands().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->getSymbolOperands());
   AffineMap layoutMap = memrefType.getLayout().getAffineMap();
   memref::AllocOp newAlloc;
   // Check if `layoutMap` is a tiled layout. Only single layout map is
-  // supported for normalizing dynamic memrefs.
+  // supported for normalizing dynamic memrefs. While normalizeMemRefType
+  // can handle normalization of memrefs with dynamic shapes, the computation
+  // of new allocation sizes is not yet supported.
   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->getAlignmentAttr());
+  } else if (newMemRefType.getNumDynamicDims() > 0 && tileSizePos.empty()) {
+    // As discussed above, handling the new size calculation for dynamically
+    // shaped memrefs during normalzation is currently unsupported.
+    return failure();  
   } else {
     newAlloc = b.create<memref::AllocOp>(allocOp->getLoc(), newMemRefType,
                                          allocOp->getAlignmentAttr());
   }
   // Replace all uses of the old memref.
   if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
                                       /*extraIndices=*/{},
                                       /*indexRemap=*/layoutMap,
LogicalResult mlir::normalizeMemRef(memref::AllocOp *allocOp) {
   assert(llvm::all_of(oldMemRef.getUsers(), [&](Operation *op) {
     return hasSingleEffect<MemoryEffects::Free>(op, oldMemRef);
   }));
   oldMemRef.replaceAllUsesWith(newAlloc);
   allocOp->erase();
   return success();
 }
 
-MemRefType mlir::normalizeMemRefType(MemRefType memrefType,
+MemRefType mlir::normalizeMemRefType(OpBuilder& builder, MemRefType memrefType,
                                      unsigned numSymbolicOperands) {
   unsigned rank = memrefType.getRank();
   if (rank == 0)
     return memrefType;
 
   if (memrefType.getLayout().isIdentity()) {
     // Either no maps is associated with this memref or this memref has
     // a trivial (identity) map.
     return memrefType;
   }
   AffineMap layoutMap = memrefType.getLayout().getAffineMap();
 
   // We don't do any checks for one-to-one'ness; we assume that it is
   // one-to-one.
 
-  // 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(layoutMap, 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();
   // FlatAffineValueConstraint may later on use symbolicOperands.
   FlatAffineValueConstraints 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.addBound(IntegerPolyhedron::LB, d, 0);
+    // The lower bound on a memref shape is always 0.
+    if (shape[d] > 0) {
+      // If the size of this dimension is statically known,
+      // add the constant upper bound.
       fac.addBound(IntegerPolyhedron::UB, d, shape[d] - 1);
     } else {
+      // Otherwise, construct an affine map to represent the dynamic
+      // bound of this dimension. The affine map constructed represents
+      // just accessing the d'th dimension of the input shape.
       memrefTypeDynDims.emplace_back(d);
+      AffineExpr ub = builder.getAffineDimExpr(d);
+      auto map = AffineMap::get(fac.getNumDimVars(), fac.getNumSymbolVars(), ub);
+      if (failed(fac.addBound(IntegerPolyhedron::UB, d, map))) {
+        return memrefType;
+      }
     }
   }
   // We compose this map with the original index (logical) space to derive
   // the upper bounds for the new index space.
   unsigned newRank = layoutMap.getNumResults();
   if (failed(fac.composeMatchingMap(layoutMap)))
     return memrefType;
   // TODO: Handle semi-affine maps.

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

for (unsigned argIndex :
     // Check whether argument is of MemRef type. Any other argument type can
     // simply be part of the final function signature.
     if (!memrefType) {
       inputTypes.push_back(argType);
       continue;
     }
     // Fetch a new memref type after normalizing the old memref to have an
     // identity map layout.
-    MemRefType newMemRefType = normalizeMemRefType(memrefType,
+    MemRefType newMemRefType = normalizeMemRefType(b, memrefType,
                                                    /*numSymbolicOperands=*/0);
+
     if (newMemRefType == memrefType || funcOp.isExternal()) {
       // Either memrefType already had an identity map or the map couldn't be
       // transformed to an identity map.
       inputTypes.push_back(newMemRefType);
       continue;
     }
 
     // Insert a new temporary argument with the new memref type.
for (unsigned resIndex :
       // Check whether result is of MemRef type. Any other argument type can
       // simply be part of the final function signature.
       if (!memrefType) {
         resultTypes.push_back(resType);
         continue;
       }
       // Computing a new memref type after normalizing the old memref to have an
       // identity map layout.
-      MemRefType newMemRefType = normalizeMemRefType(memrefType,
+      MemRefType newMemRefType = normalizeMemRefType(b, memrefType,
                                                      /*numSymbolicOperands=*/0);
       resultTypes.push_back(newMemRefType);
     }
 
     FunctionType newFuncType =
         FunctionType::get(&getContext(), /*inputs=*/inputTypes,
                           /*results=*/resultTypes);
     // Setting the new function signature for this external function.
for (unsigned resIndex : llvm::seq<unsigned>(0, oldOp->getNumResults())) {
     auto resultType = oldOp->getResult(resIndex).getType();
     MemRefType memrefType = resultType.dyn_cast<MemRefType>();
     // Check whether the operation result is MemRef type.
     if (!memrefType) {
       resultTypes.push_back(resultType);
       continue;
     }
     // Fetch a new memref type after normalizing the old memref.
-    MemRefType newMemRefType = normalizeMemRefType(memrefType,
+    MemRefType newMemRefType = normalizeMemRefType(b, memrefType,
                                                    /*numSymbolicOperands=*/0);
     if (newMemRefType == memrefType) {
       // Either memrefType already had an identity map or the map couldn't
       // be transformed to an identity map.
       resultTypes.push_back(memrefType);
       continue;
     }
     resultTypes.push_back(newMemRefType);

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

 // 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]]>) {
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x?xf32>) {
 func.func @test_norm_dynamic_not_tiled0(%arg0 : memref<1x?x?x14xf32, #map_not_tiled0>) -> () {
     %c1 = arith.constant 1 : index
     %c2 = arith.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_:%.+]] = arith.constant 1 : index
     // CHECK-DAG:       [[CST_2_:%.+]] = arith.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-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x?xf32>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x?xf32>
     // 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:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x?xf32>, 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-DAG: #[[$MAP7:.+]] = 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]]>) {
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x?x?x64xf32>) {
 func.func @test_norm_dynamic_not_tiled1(%arg0 : memref<1x?x?x14xf32, #map_not_tiled1>) -> () {
     %c1 = arith.constant 1 : index
     %c2 = arith.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_:%.+]] = arith.constant 1 : index
     // CHECK-DAG:       [[CST_2_:%.+]] = arith.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-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x?x?x64xf32>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x?x?x64xf32>
     // 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:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x?x?x64xf32>, 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-DAG: #[[$MAP8:.+]] = 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]]>) {
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x1x?x32xf32>) {
 func.func @test_norm_dynamic_not_tiled2(%arg0 : memref<1x?x?x14xf32, #map_not_tiled2>) -> () {
     %c1 = arith.constant 1 : index
     %c2 = arith.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_:%.+]] = arith.constant 1 : index
     // CHECK-DAG:       [[CST_2_:%.+]] = arith.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-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x1x?x32xf32>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x1x?x32xf32>
     // 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:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x1x?x32xf32>, 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-DAG: #[[$MAP9:.+]] = 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]]>) {
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x?x14x?x?xf32>) {
 func.func @test_norm_dynamic_not_tiled3(%arg0 : memref<1x?x?x14xf32, #map_not_tiled3>) -> () {
     %c1 = arith.constant 1 : index
     %c2 = arith.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_:%.+]] = arith.constant 1 : index
     // CHECK-DAG:       [[CST_2_:%.+]] = arith.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-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14x?x?xf32>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14x?x?xf32>
     // 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:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14x?x?xf32>, 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-DAG: #[[$MAP10:.+]] = 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]]>) {
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<1x?x1x14x32x?xf32>) {
 func.func @test_norm_dynamic_not_tiled4(%arg0 : memref<1x?x?x14xf32, #map_not_tiled4>) -> () {
     %c1 = arith.constant 1 : index
     %c2 = arith.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_:%.+]] = arith.constant 1 : index
     // CHECK-DAG:       [[CST_2_:%.+]] = arith.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-DAG:       [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x1x14x32x?xf32>
+    // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x1x14x32x?xf32>
     // 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:           "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x1x14x32x?xf32>, memref<1x?x?x14xf32, #[[$MAP9]]>) -> ()
     // CHECK:           memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP9]]>
     // CHECK:           return
 }
+
+// -----
+
+// Test that memrefs with affine maps that aren't tiled can be still be normalized when passed as function arguments.
+
+#map_dyn_arg1 = affine_map<(d0, d1, d2) -> (d1 + 3, d2 + 5)>
+#map_dyn_arg2 = affine_map<(d0, d1, d2) -> (0)>
+#map_dyn_arg3 = affine_map<(d0, d1, d2) -> (d0 * 2, d1 + d2)>
+
+// CHECK-LABEL:  func @test_func_arg_dynamic_memrefs_normalized
+// CHECK-SAME:   ([[ARG_0_:%.+]]: memref<?x?xf64>, [[ARG_1_:%.+]]: memref<1xf64>, [[ARG_2_:%.+]]: memref<?x?xf64>, [[D_0_:%.+]]: index, [[D_1_:%.+]]: index, [[D_2_:%.+]]: index) {
+func.func @test_func_arg_dynamic_memrefs_normalized(%arg0 : memref<?x?x?xf64, #map_dyn_arg1>, %arg1: memref<?x?x?xf64, #map_dyn_arg2>, %arg2 : memref<?x?x?xf64, #map_dyn_arg3>, %d0 : index, %d1 : index, %d2 : index) -> () {
+  %0 = arith.constant 0.0 : f64
+  affine.for %i = 0 to %d0 {
+    affine.for %j = 0 to %d1 {
+      affine.for %k = 0 to %d2 {
+        affine.store %0, %arg0[%i, %j, %k] : memref<?x?x?xf64, #map_dyn_arg1>
+        affine.store %0, %arg1[%i, %j, %k] : memref<?x?x?xf64, #map_dyn_arg2>
+        affine.store %0, %arg2[%i, %j, %k] : memref<?x?x?xf64, #map_dyn_arg3>
+      }
+    }
+  }
+  return
+  // CHECK: [[CST_1_:%.+]] = arith.constant 0.000000e+00 : f64
+  // CHECK: affine.for [[I_:%.+]] = 0 to [[D_0_]] {
+  // CHECK: affine.for [[J_:%.+]] = 0 to [[D_1_]] {
+  // CHECK: affine.for [[K_:%.+]] = 0 to [[D_2_]] {
+  // CHECK: affine.store [[CST_1_]], [[ARG_0_]][[[J_]] + 3, [[K_]] + 5] : memref<?x?xf64>
+  // CHECK: affine.store [[CST_1_]], [[ARG_1_]][0] : memref<1xf64>
+  // CHECK: affine.store [[CST_1_]], [[ARG_2_]][[[I_]] * 2, [[J_]] + [[K_]]] : memref<?x?xf64>
+}