NFC. Move remaining affine/memref test cases into respective dialect dirs

[bondhugula]

Summary

Move a bunch of lingering test cases from test/Transforms/ into
test/Dialect/Affine and MemRef.

mlir/test/Dialect/Affine/loop-fusion-2.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' -split-input-file | FileCheck %s
 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{fusion-maximal}))' -split-input-file | FileCheck %s --check-prefix=MAXIMAL
 
 // Part I of fusion tests in  mlir/test/Transforms/loop-fusion.mlir.
 // Part III of fusion tests in mlir/test/Transforms/loop-fusion-3.mlir
 // Part IV of fusion tests in mlir/test/Transforms/loop-fusion-4.mlir
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_at_depth_above_loop_carried_dependence(%{{.*}}: memref<64x4xf32>, %{{.*}}: memref<64x4xf32>) {
 func.func @should_fuse_at_depth_above_loop_carried_dependence(%arg0: memref<64x4xf32>, %arg1: memref<64x4xf32>) {
   %out = memref.alloc() : memref<64x4xf32>
   %0 = arith.constant 0.0 : f32
   affine.for %i0 = 0 to 64 {
     affine.for %i1 = 0 to 4 {
       affine.store %0, %out[%i0, %i1] : memref<64x4xf32>
     }
   }
   affine.for %i2 = 0 to 4 {
     affine.for %i3 = 0 to 4 {
       affine.for %i4 = 0 to 16 {
         %v = affine.load %arg1[16 * %i3 - %i4 + 15, %i2] : memref<64x4xf32>
         "op0"(%v) : (f32) -> ()
       }
       affine.for %i5 = 0 to 4 {
         affine.for %i6 = 0 to 16 {
           %v = affine.load %arg0[16 * %i5 - %i6 + 15, %i3] : memref<64x4xf32>
           "op1"(%v) : (f32) -> ()
         }
         affine.for %i7 = 0 to 16 {
           %r = "op2"() : () -> (f32)
           %v = affine.load %out[16 * %i5 + %i7, %i2] : memref<64x4xf32>
           %s = arith.addf %v, %r : f32
           affine.store %s, %out[16 * %i5 + %i7, %i2] : memref<64x4xf32>
         }
       }
     }
   }
 
   // We can fuse source loop nest '%i0' into dst loop nest '%i2', but the
   // depth at which we can insert the src loop nest slice into the dst loop
   // lest must be decreased because of a loop carried dependence on loop '%i3'.
   // As a result, the source loop nest is inserted at dst loop nest depth 1,
   // just above the loop with the carried dependence. In addition, the source
   // loop nest iteration bounds on its loop '%i1' are reduced to 1, so the
   // memref size can be reduced to 128x1xf32.
 
   // CHECK:       memref.alloc() : memref<64x1xf32>
   // CHECK:       affine.for %{{.*}} = 0 to 4 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 64 {
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}, 0] : memref<64x1xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 4 {
   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 16 {
   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, %{{.*}}] : memref<64x4xf32>
   // CHECK-NEXT:        "op0"(%{{.*}}) : (f32) -> ()
   // CHECK-NEXT:      }
   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 4 {
   // CHECK-NEXT:        affine.for %{{.*}} = 0 to 16 {
   // CHECK-NEXT:          affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, %{{.*}}] : memref<64x4xf32>
   // CHECK-NEXT:          "op1"(%{{.*}}) : (f32) -> ()
   // CHECK-NEXT:        }
   // CHECK-NEXT:        affine.for %{{.*}} = 0 to 16 {
   // CHECK-NEXT:          %{{.*}} = "op2"() : () -> f32
   // CHECK:               affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, 0] : memref<64x1xf32>
   // CHECK-NEXT:          arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK:               affine.store %{{.*}}, %{{.*}}[%{{.*}} * 16 + %{{.*}}, 0] : memref<64x1xf32>
   // CHECK-NEXT:        }
   // CHECK-NEXT:      }
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_only_two_loops_and_remove_producer() {
 func.func @should_fuse_only_two_loops_and_remove_producer() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %a[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %a[%i1] : memref<10xf32>
     affine.store %v0, %b[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     %v1 = affine.load %a[%i2] : memref<10xf32>
     affine.store %v1, %b[%i2] : memref<10xf32>
   }
 
   // On the first visit to '%i2', the fusion algorithm can not fuse loop nest
   // '%i0' into '%i2' because of the dependences '%i0' and '%i2' each have on
   // '%i1'. Then, '%i0' is fused into '%i1' and no private memref is created for
   // memref '%a' to be able to remove '%i0' and still preserve the depencence on
   // '%a' with '%i2'.
   // TODO: Alternatively, we could fuse '%i0' into '%i1' with a private memref,
   // the dependence between '%i0' and '%i1' on memref '%a' would no longer exist,
   // and '%i0' could be fused into '%i2' as well. Note that this approach would
   // duplicate the computation in loop nest '%i0' to loop nests '%i1' and '%i2',
   // which would limit its profitability.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:   return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_after_one_loop_interchange() {
 func.func @should_fuse_after_one_loop_interchange() {
   %a = memref.alloc() : memref<10xf32>
 
   %cf0 = arith.constant 0.0 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cf0, %a[%i0] : memref<10xf32>
   }
 
   affine.for %i1 = 0 to 5 {
     affine.for %i2 = 0 to 10 {
       %v0 = affine.load %a[%i2] : memref<10xf32>
       affine.store %v0, %a[%i2] : memref<10xf32>
     }
   }
 
   // The dependence between the load and affine.store is carried on loop '%i1', and
   // cannot be fused with loop '%i0' without violating this dependence.
   // Once loops '%i1' and %i2' are interchanged, loop '%i0' can be fused
   // at loop depth 1, because the loop carrying the dependence has been
   // interchanged and is now at depth 2.
 
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 5 {
   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 // CHECK-LABEL: func @should_fuse_after_two_loop_interchanges() {
 func.func @should_fuse_after_two_loop_interchanges() {
   %a = memref.alloc() : memref<6x8xf32>
 
   %cf0 = arith.constant 0.0 : f32
   affine.for %i0 = 0 to 6 {
     affine.for %i1 = 0 to 8 {
       affine.store %cf0, %a[%i0, %i1] : memref<6x8xf32>
     }
   }
 
   affine.for %i2 = 0 to 4 {
     affine.for %i3 = 0 to 6 {
       affine.for %i4 = 0 to 2 {
         affine.for %i5 = 0 to 8 {
           %v0 = affine.load %a[%i3, %i5] : memref<6x8xf32>
           %v1 = arith.addf %v0, %v0 : f32
           affine.store %v1, %a[%i3, %i5] : memref<6x8xf32>
         }
       }
     }
   }
 
   // The dependence between the load and affine.store is carried on loops '%i2' and
   // '%i4', and cannot be fused with loop '%i0' without violating this
   // dependence.
   // Once loop '%i2' is interchanged with loop '%i3', and again with loop
   // '%i5', then loop '%i0' can be fused at loop depth 2, because the loop
   // carrying the dependences have been interchanged with loops at depth > 2.
 
   // CHECK:       affine.for %{{.*}} = 0 to 6 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 8 {
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 4 {
   // CHECK-NEXT:        affine.for %{{.*}} = 0 to 2 {
   // CHECK-NEXT:          affine.load %{{.*}}[0, 0] : memref<1x1xf32>
   // CHECK-NEXT:          arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:          affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
   // CHECK-NEXT:        }
   // CHECK-NEXT:      }
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 func.func @should_fuse_live_out_writer(%arg0 : memref<10xf32>) -> memref<10xf32> {
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %arg0[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %1 = affine.load %arg0[%i1] : memref<10xf32>
     affine.store %1, %arg0[%i1] : memref<10xf32>
   }
   return %arg0 : memref<10xf32>
 
   // CHECK:       %{{.*}} = arith.constant 0.000000e+00 : f32
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return %{{.*}} : memref<10xf32>
 }
 
 // -----
 
 // The fused slice has 16 iterations from along %i0.
 
 // CHECK-DAG: [[$MAP_LB:#map[0-9]*]] = affine_map<(d0) -> (d0 * 16)>
 // CHECK-DAG: [[$MAP_UB:#map[0-9]*]] = affine_map<(d0) -> (d0 * 16 + 16)>
 
 // CHECK-LABEL: slice_tile
 func.func @slice_tile(%arg0: memref<128x8xf32>, %arg1: memref<32x8xf32>, %0 : f32) -> memref<32x8xf32> {
   affine.for %i0 = 0 to 32 {
     affine.for %i1 = 0 to 8 {
       affine.store %0, %arg1[%i0, %i1] : memref<32x8xf32>
     }
   }
   affine.for %i = 0 to 2 {
     affine.for %j = 0 to 8 {
       affine.for %k = 0 to 8 {
         affine.for %kk = 0 to 16 {
           %v = affine.load %arg0[16 * %k + %kk, %j] : memref<128x8xf32>
           %r = "foo"(%v) : (f32) -> f32
         }
         affine.for %ii = 0 to 16 {
           %v = affine.load %arg1[16 * %i + %ii, %j] : memref<32x8xf32>
           %s = arith.addf %v, %v : f32
           affine.store %s, %arg1[16 * %i + %ii, %j] : memref<32x8xf32>
         }
       }
     }
   }
   return %arg1 : memref<32x8xf32>
 }
 // CHECK:       affine.for %{{.*}} = 0 to 2 {
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 8 {
 // CHECK-NEXT:      affine.for %{{.*}} = [[$MAP_LB]](%{{.*}}) to [[$MAP_UB]](%{{.*}}) {
 // CHECK-NEXT:        affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<32x8xf32>
 // CHECK-NEXT:      }
 // CHECK-NEXT:      affine.for %{{.*}} = 0 to 8 {
 // CHECK-NEXT:        affine.for %{{.*}} = 0 to 16 {
 // CHECK-NEXT:          affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<128x8xf32>
 // CHECK-NEXT:          "foo"(%{{.*}}) : (f32) -> f32
 // CHECK-NEXT:        }
 // CHECK-NEXT:        affine.for %{{.*}} = 0 to 16 {
 // CHECK-NEXT:          affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<32x8xf32>
 // CHECK-NEXT:          arith.addf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:          affine.store %{{.*}}, %{{.*}}[%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<32x8xf32>
 // CHECK-NEXT:        }
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return %{{.*}} : memref<32x8xf32>
 // CHECK-NEXT:}
 
 // -----
 
 // Test case which illustrates fix for b/126454413
 func.func @test_add_slice_bounds() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
   %c0 = arith.constant 0 : index
 
   affine.for %i0 = 0 to 10 {
     affine.for %i1 = 0 to 10 {
       affine.for %i2 = 0 to 10 {
         %a0 = affine.apply affine_map<(d0) -> (d0)> (%i0)
         %a1 = affine.apply affine_map<(d0) -> (d0)> (%i0)
         %a2 = affine.apply affine_map<(d0, d1) -> (d0 - d1)> (%a0, %a1)
         affine.store %cf7, %a[%a2] : memref<10xf32>
       }
     }
   }
   affine.for %i3 = 0 to 10 {
     affine.for %i4 = 0 to 10 {
       affine.for %i5 = 0 to 10 {
         %v0 = affine.load %a[%c0] : memref<10xf32>
       }
     }
   }
 
 // CHECK:        affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:       affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:         affine.apply #map(%{{.*}})
 // CHECK-NEXT:         affine.apply #map(%{{.*}})
 // CHECK-NEXT:         affine.apply #map1(%{{.*}}, %{{.*}})
 // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:       }
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
 // CHECK-NEXT:   affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:       affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:       }
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
   return
 }
 
 // -----
 
 func.func @should_fuse_init_loops_siblings_then_shared_producer(%arg0: memref<10x10xf32>, %arg1: memref<10x10xf32>) {
   %0 = memref.alloc() : memref<10x10xf32>
   %cst = arith.constant 0.000000e+00 : f32
   %cst_0 = arith.constant 1.000000e+00 : f32
   %cst_1 = arith.constant 7.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.for %i1 = 0 to 10 {
       affine.store %cst_1, %0[%i0, %i1] : memref<10x10xf32>
     }
   }
   affine.for %i2 = 0 to 3 {
     affine.for %i3 = 0 to 3 {
       affine.store %cst, %arg0[%i2, %i3] : memref<10x10xf32>
     }
   }
   affine.for %i4 = 0 to 3 {
     affine.for %i5 = 0 to 3 {
       %1 = affine.load %0[%i4, %i5] : memref<10x10xf32>
       %2 = affine.load %arg0[%i4, %i5] : memref<10x10xf32>
       %3 = arith.mulf %1, %2 : f32
       affine.store %3, %arg0[%i4, %i5] : memref<10x10xf32>
     }
   }
   affine.for %i6 = 0 to 3 {
     affine.for %i7 = 0 to 3 {
       affine.store %cst_0, %arg1[%i6, %i7] : memref<10x10xf32>
     }
   }
   affine.for %i8 = 0 to 3 {
     affine.for %i9 = 0 to 3 {
       %4 = affine.load %0[%i8, %i9] : memref<10x10xf32>
       %5 = affine.load %arg1[%i8, %i9] : memref<10x10xf32>
       %6 = arith.addf %4, %5 : f32
       affine.store %6, %arg1[%i8, %i9] : memref<10x10xf32>
     }
   }
 
   // Pass 1: should fuse single-use producer loop nests into their unique user,
   //         so '%i2' will fuse into '%i4' and '%i6' will fuse into '%i8'.
   // Pass 2: should fuse sibling loop nests which share no dependence edges,
   //         so should fuse '%i4' into '%i8'.
   // Pass 3: should fuse single-use producer loop nest '%i0' into '%i8'. Note
   //         that loop nest '%i0' now has a single user after Pass 2 fused its
   //         two users together).
 
 // CHECK:        affine.for %{{.*}} = 0 to 3 {
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 3 {
 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:       affine.load %{{.*}}[0, 0] : memref<1x1xf32>
 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:       arith.mulf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:       affine.load %{{.*}}[0, 0] : memref<1x1xf32>
 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:       arith.addf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
 // CHECK-NEXT:   return
 
   return
 }
 
 // -----
 
 func.func @two_matrix_vector_products() {
   %in_matrix = memref.alloc() : memref<10x10xf32>
   %in_vec0 = memref.alloc() : memref<10xf32>
   %in_vec1 = memref.alloc() : memref<10xf32>
   %out_vec0 = memref.alloc() : memref<10xf32>
   %out_vec1 = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   // Populate input matrix.
   affine.for %i0 = 0 to 10 {
     affine.for %i1 = 0 to 10 {
       affine.store %cf7, %in_matrix[%i0, %i1] : memref<10x10xf32>
     }
   }
   // out_vec0 = in_matrix x in_vec0
   affine.for %i2 = 0 to 10 {
     affine.for %i3 = 0 to 10 {
       %v0 = affine.load %in_matrix[%i2, %i3] : memref<10x10xf32>
       %v1 = affine.load %in_vec0[%i3] : memref<10xf32>
       %v2 = arith.mulf %v0, %v1 : f32
       %v3 = affine.load %out_vec0[%i3] : memref<10xf32>
       %v4 = arith.addf %v2, %v3 : f32
       affine.store %v4, %out_vec0[%i3] : memref<10xf32>
     }
   }
   // out_vec1 = in_matrix x in_vec1
   affine.for %i4 = 0 to 10 {
     affine.for %i5 = 0 to 10 {
       %v5 = affine.load %in_matrix[%i4, %i5] : memref<10x10xf32>
       %v6 = affine.load %in_vec1[%i5] : memref<10xf32>
       %v7 = arith.mulf %v5, %v6 : f32
       %v8 = affine.load %out_vec1[%i5] : memref<10xf32>
       %v9 = arith.addf %v7, %v8 : f32
       affine.store %v9, %out_vec1[%i5] : memref<10xf32>
     }
   }
 
 // CHECK:        affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, 0] : memref<10x1xf32>
 // CHECK-NEXT:     }
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}, 0] : memref<10x1xf32>
 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:       arith.mulf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:       arith.addf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:     }
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}, 0] : memref<10x1xf32>
 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:       arith.mulf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:       affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:       arith.addf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
 // CHECK-NEXT:   return
   return
 }
 
 // -----
 
 func.func @should_not_slice_past_slice_barrier() {
   %0 = memref.alloc() : memref<100x16xf32>
   affine.for %i0 = 0 to 100 {
     affine.for %i1 = 0 to 16 {
       %1 = "op1"() : () -> f32
       affine.store %1, %0[%i0, %i1] : memref<100x16xf32>
     } {slice_fusion_barrier = true}
   }
   affine.for %i2 = 0 to 100 {
     affine.for %i3 = 0 to 16 {
       %2 = affine.load %0[%i2, %i3] : memref<100x16xf32>
       "op2"(%2) : (f32) -> ()
     }
   }
   // The 'slice_fusion_barrier' attribute on '%i1' prevents slicing the
   // iteration space of '%i1' and any enclosing loop nests.
 // CHECK:        affine.for %{{.*}} = 0 to 100 {
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 16 {
 // CHECK-NEXT:       %{{.*}} = "op1"() : () -> f32
 // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
 // CHECK-NEXT:     } {slice_fusion_barrier = true}
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 16 {
 // CHECK-NEXT:       affine.load %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
 // CHECK-NEXT:       "op2"(%{{.*}}) : (f32) -> ()
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
   return
 }
 
 // -----
 
 #map = affine_map<(d0, d1) -> (d0 * 16 + d1)>
 func.func @fuse_across_dim_mismatch(%arg0: memref<4x4x16x1xf32>, %arg1: memref<144x9xf32>, %arg2: memref<9xf32>) {
   %1 = memref.alloc() : memref<144x4xf32>
   %2 = arith.constant 0.0 : f32
   affine.for %i2 = 0 to 9 {
     affine.for %i3 = 0 to 4 {
       affine.for %i5 = 0 to 16 {
         %7 = affine.apply #map(%i2, %i5)
         affine.store %2, %1[%7, %i3] : memref<144x4xf32>
       }
     }
   }
   affine.for %i6 = 0 to 9 {
     affine.for %i7 = 0 to 9 {
       affine.for %i8 = 0 to 4 {
         affine.for %i10 = 0 to 16 {
           %10 = affine.apply #map(%i6, %i10)
           %11 = affine.load %1[%10, %i8] : memref<144x4xf32>
         }
       }
     }
   }
   return
 }
 // MAXIMAL:      #[[$MAP:.*]] = affine_map<(d0, d1) -> (d0 * 16 + d1)>
 // MAXIMAL-LABEL: func @fuse_across_dim_mismatch
 // MAXIMAL:        memref.alloc() : memref<1x1xf32>
 // MAXIMAL:        affine.for %{{.*}} = 0 to 9 {
 // MAXIMAL-NEXT:    affine.for %{{.*}} = 0 to 9 {
 // MAXIMAL-NEXT:      affine.for %{{.*}} = 0 to 4 {
 // MAXIMAL-NEXT:        affine.for %{{.*}} = 0 to 16 {
 // MAXIMAL-NEXT:          affine.apply #[[$MAP]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:          affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
 // MAXIMAL-NEXT:          affine.apply #[[$MAP]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:          affine.load %{{.*}}[0, 0] : memref<1x1xf32>
 // MAXIMAL-NEXT:        }
 // MAXIMAL-NEXT:      }
 // MAXIMAL-NEXT:    }
 // MAXIMAL-NEXT:  }
 
 // -----
 
 #map3 = affine_map<(d0, d1) -> ((d0 * 72 + d1) floordiv 2304)>
 #map4 = affine_map<(d0, d1) -> (((d0 * 72 + d1) mod 2304) floordiv 1152)>
 #map5 = affine_map<(d0, d1) -> (((((d0 * 72 + d1) mod 2304) mod 1152) floordiv 9) floordiv 8)>
 #map6 = affine_map<(d0, d1) -> (((((d0 * 72 + d1) mod 2304) mod 1152) mod 9) floordiv 3)>
 #map7 = affine_map<(d0, d1) -> (((((d0 * 72 + d1) mod 2304) mod 1152) mod 9) mod 3)>
 #map10 = affine_map<(d0, d1) -> (d0 * 16 + d1)>
 #map11 = affine_map<(d0, d1) -> (d0 * 16 + d1)>
 #map12 = affine_map<(d0, d1) -> (d0 * 16 - d1 + 15)>
 func.func @fuse_across_varying_dims_complex(%arg0: f32) {
   %c0 = arith.constant 0 : index
   %0 = memref.alloc() : memref<2x2x3x3x16x1xf32>
   %1 = memref.alloc() : memref<64x9xf32>
   %2 = memref.alloc() : memref<144x4xf32>
   affine.for %i0 = 0 to 64 {
     affine.for %i1 = 0 to 9 {
       %4 = affine.apply #map3(%i0, %i1)
       %5 = affine.apply #map4(%i0, %i1)
       %6 = affine.apply #map5(%i0, %i1)
       %7 = affine.apply #map6(%i0, %i1)
       %8 = affine.apply #map7(%i0, %i1)
       %9 = affine.load %0[%4, %5, %7, %8, %6, %c0] : memref<2x2x3x3x16x1xf32>
       affine.store %9, %1[%i0, %i1] : memref<64x9xf32>
     }
   }
   affine.for %i2 = 0 to 9 {
     affine.for %i3 = 0 to 4 {
       affine.for %i4 = 0 to 16 {
         %10 = affine.apply #map10(%i3, %i4)
         %11 = affine.load %1[%10, %i2] : memref<64x9xf32>
       }
       affine.for %i5 = 0 to 16 {
         %14 = affine.apply #map11(%i2, %i5)
         affine.store %arg0, %2[%14, %i3] : memref<144x4xf32>
       }
     }
   }
   affine.for %i6 = 0 to 9 {
     affine.for %i7 = 0 to 9 {
       affine.for %i8 = 0 to 4 {
         affine.for %i9 = 0 to 16 {
           %15 = affine.apply #map12(%i8, %i9)
           %16 = affine.load %1[%15, %i7] : memref<64x9xf32>
         }
       }
     }
   }
   return
 }
 // MAXIMAL-DAG: [[$MAP0:#map[0-9]*]] = affine_map<(d0, d1) -> ((d0 * 72 + d1) floordiv 2304)>
 // MAXIMAL-DAG: [[$MAP1:#map[0-9]*]] = affine_map<(d0, d1) -> (((d0 * 72 + d1) mod 2304) floordiv 1152)>
 // MAXIMAL-DAG: [[$MAP2:#map[0-9]*]] = affine_map<(d0, d1) -> ((((d0 * 72 + d1) mod 1152) floordiv 9) floordiv 8)>
 // MAXIMAL-DAG: [[$MAP3:#map[0-9]*]] = affine_map<(d0, d1) -> ((d1 mod 9) floordiv 3)>
 // MAXIMAL-DAG: [[$MAP4:#map[0-9]*]] = affine_map<(d0, d1) -> (d1 mod 3)>
 // MAXIMAL-DAG: [[$MAP7:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 * 16 + d1)>
 // MAXIMAL-DAG: [[$MAP8:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 * 16 - d1 + 15)>
 // MAXIMAL-LABEL: func @fuse_across_varying_dims_complex
 // MAXIMAL-NEXT:  memref.alloc() : memref<64x1xf32>
 // MAXIMAL-NEXT:  arith.constant 0 : index
 // MAXIMAL-NEXT:  memref.alloc() : memref<2x2x3x3x16x1xf32>
 // MAXIMAL-NEXT:  memref.alloc() : memref<144x4xf32>
 // MAXIMAL-NEXT:  affine.for %{{.*}} = 0 to 9 {
 // MAXIMAL-NEXT:    affine.for %{{.*}} = 0 to 4 {
 // MAXIMAL-NEXT:      affine.for %{{.*}} = 0 to 16 {
 // MAXIMAL-NEXT:        affine.for %{{.*}} = 0 to 64 {
 // MAXIMAL-NEXT:          affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:          affine.apply [[$MAP1]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:          affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:          affine.apply [[$MAP3]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:          affine.apply [[$MAP4]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:          affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<2x2x3x3x16x1xf32>
 // MAXIMAL-NEXT:          affine.store %{{.*}}, %{{.*}}[%{{.*}}, 0] : memref<64x1xf32>
 // MAXIMAL-NEXT:        }
 // MAXIMAL-NEXT:        affine.apply [[$MAP7]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:        affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, 0] : memref<64x1xf32>
 // MAXIMAL-NEXT:        affine.for %{{.*}} = 0 to 9 {
 // MAXIMAL-NEXT:          affine.for %{{.*}} = 0 to 4 {
 // MAXIMAL-NEXT:            affine.for %{{.*}} = 0 to 16 {
 // MAXIMAL-NEXT:              affine.apply [[$MAP8]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:              affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, 0] : memref<64x1xf32>
 // MAXIMAL-NEXT:            }
 // MAXIMAL-NEXT:          }
 // MAXIMAL-NEXT:        }
 // MAXIMAL-NEXT:      }
 // MAXIMAL-NEXT:      affine.for %{{.*}} = 0 to 16 {
 // MAXIMAL-NEXT:        affine.apply [[$MAP7]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:        affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<144x4xf32>
 // MAXIMAL-NEXT:      }
 // MAXIMAL-NEXT:    }
 // MAXIMAL-NEXT:  }
 
 // -----
 
 func.func @should_fuse_with_slice_union() {
   %a = memref.alloc() : memref<100xf32>
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
 
   affine.for %i0 = 0 to 100 {
     affine.store %cf0, %a[%i0]: memref<100xf32>
   }
 
   affine.for %i1 = 10 to 20 {
     %v0 = affine.load %a[%i1]: memref<100xf32>
     affine.for %i2 = 15 to 25 {
       %v1 = affine.load %a[%i2]: memref<100xf32>
     }
   }
   // The union of two slice bounds (calculated between the store and each of
   // the loads) is computed and used in the fusion cost calculation, index
   // remapping, and private memref size. The result is that the temporary
   // memref is reduced from 100xf32 to 15xf32 and properly indexed by
   // the fused loops based on the union calculation.
 // CHECK:      affine.for %{{.*}} = 10 to 20 {
 // CHECK-NEXT:   affine.for %{{.*}} = 10 to 25 {
 // CHECK-NEXT:     affine.store %{{.*}}, %{{.*}}[%{{.*}} - 10] : memref<15xf32>
 // CHECK-NEXT:   }
 // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}} - 10] : memref<15xf32>
 // CHECK-NEXT:   affine.for %{{.*}} = 15 to 25 {
 // CHECK-NEXT:     affine.load %{{.*}}[%{{.*}} - 10] : memref<15xf32>
 // CHECK-NEXT:   }
 // CHECK-NEXT: }
 // CHECK-NEXT: return
   return
 }
 
 // -----
 
 func.func @affine_add_mm_fused(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>, %arg2: memref<1024x1024xf32>, %arg3: memref<1024x1024xf32>) {
   affine.for %i2 = 0 to 1024 {
     affine.for %i3 = 0 to 1024 {
       %0 = affine.load %arg3[%i2, %i3] : memref<1024x1024xf32>
       %1 = affine.load %arg2[%i2, %i3] : memref<1024x1024xf32>
       %2 = arith.addf %1, %0 : f32
       affine.store %2, %arg2[%i2, %i3] : memref<1024x1024xf32>
     }
   }
   affine.for %i4 = 0 to 1024 {
     affine.for %i5 = 0 to 1024 {
       affine.for %i6 = 0 to 1024 {
         %3 = affine.load %arg1[%i6, %i5] : memref<1024x1024xf32>
         %4 = affine.load %arg0[%i4, %i6] : memref<1024x1024xf32>
         %5 = arith.mulf %4, %3 : f32
         %6 = affine.load %arg2[%i4, %i5] : memref<1024x1024xf32>
         %7 = arith.addf %6, %5 : f32
         affine.store %7, %arg2[%i4, %i5] : memref<1024x1024xf32>
       }
     }
   }
   // Should fuse elementwise add loop at loop depth 2, above loop-carried
   // dependence between load/store on '%arg2', carried on reduction loop %i6.
   // CHECK:       affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:      arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:        arith.mulf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:        arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:        affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:      }
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   return
 }
 
 // -----
 
 func.func @affine_2mm_fused(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>, %arg2: memref<1024x1024xf32>, %arg3: memref<1024x1024xf32>, %arg4: memref<1024x1024xf32>) {
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 1024 {
     affine.for %i1 = 0 to 1024 {
       affine.store %cst, %arg2[%i0, %i1] : memref<1024x1024xf32>
     }
   }
   affine.for %i2 = 0 to 1024 {
     affine.for %i3 = 0 to 1024 {
       affine.store %cst, %arg4[%i2, %i3] : memref<1024x1024xf32>
     }
   }
   affine.for %i4 = 0 to 1024 {
     affine.for %i5 = 0 to 1024 {
       affine.for %i6 = 0 to 1024 {
         %0 = affine.load %arg1[%i6, %i5] : memref<1024x1024xf32>
         %1 = affine.load %arg0[%i4, %i6] : memref<1024x1024xf32>
         %2 = arith.mulf %1, %0 : f32
         %3 = affine.load %arg2[%i4, %i5] : memref<1024x1024xf32>
         %4 = arith.addf %3, %2 : f32
         affine.store %4, %arg2[%i4, %i5] : memref<1024x1024xf32>
       }
     }
   }
   affine.for %i7 = 0 to 1024 {
     affine.for %i8 = 0 to 1024 {
       affine.for %i9 = 0 to 1024 {
         %5 = affine.load %arg1[%i9, %i8] : memref<1024x1024xf32>
         %6 = affine.load %arg0[%i7, %i9] : memref<1024x1024xf32>
         %7 = arith.mulf %6, %5 : f32
         %8 = affine.load %arg4[%i7, %i8] : memref<1024x1024xf32>
         %9 = arith.addf %8, %7 : f32
         affine.store %9, %arg4[%i7, %i8] : memref<1024x1024xf32>
       }
     }
   }
 
   // Should fuse MM initialization loops into their consumers, then fuse the
   // two matmul loops together for input reuse on '%arg0/%arg1'.
 
   // CHECK:        affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:     affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         arith.mulf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:       }
   // CHECK-NEXT:     }
   // CHECK-NEXT:     affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         arith.mulf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:       }
   // CHECK-NEXT:     }
   // CHECK-NEXT:   }
 
   return
 }
 
 // -----
 
 func.func @affine_2_dependent_mm_fused(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>, %arg2: memref<1024x1024xf32>, %arg3: memref<1024x1024xf32>, %arg4: memref<1024x1024xf32>) {
   affine.for %i0 = 0 to 1024 {
     affine.for %i1 = 0 to 1024 {
       affine.for %i2 = 0 to 1024 {
         %0 = affine.load %arg1[%i2, %i1] : memref<1024x1024xf32>
         %1 = affine.load %arg0[%i0, %i2] : memref<1024x1024xf32>
         %2 = arith.mulf %1, %0 : f32
         %3 = affine.load %arg2[%i0, %i1] : memref<1024x1024xf32>
         %4 = arith.addf %3, %2 : f32
         affine.store %4, %arg2[%i0, %i1] : memref<1024x1024xf32>
       }
     }
   }
   affine.for %i3 = 0 to 1024 {
     affine.for %i4 = 0 to 1024 {
       affine.for %i5 = 0 to 1024 {
         %5 = affine.load %arg3[%i5, %i4] : memref<1024x1024xf32>
         %6 = affine.load %arg2[%i3, %i5] : memref<1024x1024xf32>
         %7 = arith.mulf %6, %5 : f32
         %8 = affine.load %arg4[%i3, %i4] : memref<1024x1024xf32>
         %9 = arith.addf %8, %7 : f32
         affine.store %9, %arg4[%i3, %i4] : memref<1024x1024xf32>
       }
     }
   }
 
   // CHECK:        affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:     affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         arith.mulf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:       }
   // CHECK-NEXT:     }
   // CHECK-NEXT:     affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 {
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         arith.mulf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:         arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>
   // CHECK-NEXT:       }
   // CHECK-NEXT:     }
   // CHECK-NEXT:   }
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_self_dependence_multi_store_producer() {
 func.func @should_fuse_self_dependence_multi_store_producer() {
   %m = memref.alloc() : memref<10xf32>
   %local_m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %local_m[%i0] : memref<10xf32>
     %v0 = affine.load %local_m[%i0] : memref<10xf32>
     affine.store %v0, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v1 = affine.load %m[%i1] : memref<10xf32>
   }
   // CHECK:      affine.for %[[i0:.*]] = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, [[LOCAL_M:%.*]][%[[i0]]] : memref<10xf32>
   // CHECK-NEXT:   [[v0:%.*]] = affine.load [[LOCAL_M]][%[[i0]]] : memref<10xf32>
   // CHECK-NEXT:   affine.store [[v0]], %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_dead_multi_store_producer() {
 func.func @should_fuse_dead_multi_store_producer() {
   %m = memref.alloc() : memref<10xf32>
   %dead_m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %dead_m[%i0] : memref<10xf32>
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // CHECK:      affine.for %[[i0:.*]] = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%[[i0]]] : memref<10xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_function_live_out_multi_store_producer
 func.func @should_fuse_function_live_out_multi_store_producer(%live_in_out_m : memref<10xf32>) {
   %m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %live_in_out_m[%i0] : memref<10xf32>
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // CHECK:      affine.for %[[i0:.*]] = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%[[i0]]] : memref<10xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // Add further tests in mlir/test/Transforms/loop-fusion-4.mlir

mlir/test/Dialect/Affine/loop-fusion-3.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' -split-input-file | FileCheck %s
 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{fusion-maximal}))' -split-input-file | FileCheck %s --check-prefix=MAXIMAL
 
 // Part I of fusion tests in  mlir/test/Transforms/loop-fusion.mlir.
 // Part II of fusion tests in mlir/test/Transforms/loop-fusion-2.mlir
 // Part IV of fusion tests in mlir/test/Transforms/loop-fusion-4.mlir
 
 // -----
 
 // Test case from github bug 777.
 // CHECK-LABEL: func @mul_add_0
 func.func @mul_add_0(%arg0: memref<3x4xf32>, %arg1: memref<4x3xf32>, %arg2: memref<3x3xf32>, %arg3: memref<3x3xf32>) {
   %cst = arith.constant 0.000000e+00 : f32
   %0 = memref.alloc() : memref<3x3xf32>
   affine.for %arg4 = 0 to 3 {
     affine.for %arg5 = 0 to 3 {
       affine.store %cst, %0[%arg4, %arg5] : memref<3x3xf32>
     }
   }
   affine.for %arg4 = 0 to 3 {
     affine.for %arg5 = 0 to 3 {
       affine.for %arg6 = 0 to 4 {
         %1 = affine.load %arg1[%arg6, %arg5] : memref<4x3xf32>
         %2 = affine.load %arg0[%arg4, %arg6] : memref<3x4xf32>
         %3 = arith.mulf %2, %1 : f32
         %4 = affine.load %0[%arg4, %arg5] : memref<3x3xf32>
         %5 = arith.addf %4, %3 : f32
         affine.store %5, %0[%arg4, %arg5] : memref<3x3xf32>
       }
     }
   }
   affine.for %arg4 = 0 to 3 {
     affine.for %arg5 = 0 to 3 {
       %6 = affine.load %arg2[%arg4, %arg5] : memref<3x3xf32>
       %7 = affine.load %0[%arg4, %arg5] : memref<3x3xf32>
       %8 = arith.addf %7, %6 : f32
       affine.store %8, %arg3[%arg4, %arg5] : memref<3x3xf32>
     }
   }
   // CHECK:      affine.for %[[i0:.*]] = 0 to 3 {
   // CHECK-NEXT:   affine.for %[[i1:.*]] = 0 to 3 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
   // CHECK-NEXT:     affine.for %[[i2:.*]] = 0 to 4 {
   // CHECK-NEXT:       affine.load %{{.*}}[%[[i2]], %[[i1]]] : memref<4x3xf32>
   // CHECK-NEXT:       affine.load %{{.*}}[%[[i0]], %[[i2]]] : memref<3x4xf32>
   // CHECK-NEXT:       arith.mulf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:       affine.load %{{.*}}[0, 0] : memref<1x1xf32>
   // CHECK-NEXT:       arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:       affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
   // CHECK-NEXT:     }
   // CHECK-NEXT:     affine.load %{{.*}}[%[[i0]], %[[i1]]] : memref<3x3xf32>
   // CHECK-NEXT:     affine.load %{{.*}}[0, 0] : memref<1x1xf32>
   // CHECK-NEXT:     arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:     affine.store %{{.*}}, %{{.*}}[%[[i0]], %[[i1]]] : memref<3x3xf32>
   // CHECK-NEXT:   }
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // Verify that 'fuseProducerConsumerNodes' fuse a producer loop with a store
 // that has multiple outgoing edges.
 
 // CHECK-LABEL: func @should_fuse_multi_outgoing_edge_store_producer
 func.func @should_fuse_multi_outgoing_edge_store_producer(%a : memref<1xf32>) {
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %arg0 = 0 to 1 {
     affine.store %cst, %a[%arg0] : memref<1xf32>
   }
 
   affine.for %arg0 = 0 to 1 {
     %0 = affine.load %a[%arg0] : memref<1xf32>
   }
 
   affine.for %arg0 = 0 to 1 {
     %0 = affine.load %a[%arg0] : memref<1xf32>
   }
   // CHECK:      affine.for %{{.*}} = 0 to 1 {
   // CHECK-NEXT:   affine.store
   // CHECK-NEXT:   affine.load
   // CHECK-NEXT:   affine.load
   // CHECK-NEXT: }
 
   return
 }
 
 // -----
 
 // Verify that 'fuseProducerConsumerNodes' fuses a producer loop that: 1) has
 // multiple outgoing edges, 2) producer store has a single outgoing edge.
 // Sibling loop fusion should not fuse any of these loops due to
 // dependencies on external memrefs '%a' and '%b'.
 
 // CHECK-LABEL: func @should_fuse_producer_with_multi_outgoing_edges
 func.func @should_fuse_producer_with_multi_outgoing_edges(%a : memref<1xf32>, %b : memref<1xf32>) {
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %arg0 = 0 to 1 {
     %0 = affine.load %a[%arg0] : memref<1xf32>
     affine.store %cst, %b[%arg0] : memref<1xf32>
   }
 
   affine.for %arg0 = 0 to 1 {
     affine.store %cst, %a[%arg0] : memref<1xf32>
     %1 = affine.load %b[%arg0] : memref<1xf32>
   }
   // CHECK: affine.for %{{.*}} = 0 to 1
   // CHECK-NEXT: affine.load %[[A:.*]][{{.*}}]
   // CHECK-NEXT: affine.store %{{.*}}, %[[B:.*]][{{.*}}]
   // CHECK-NEXT: affine.store %{{.*}}, %[[A]]
   // CHECK-NEXT: affine.load %[[B]]
   // CHECK-NOT: affine.for %{{.*}}
   // CHECK: return
   return
 }
 
 // MAXIMAL-LABEL: func @reshape_into_matmul
 func.func @reshape_into_matmul(%lhs : memref<1024x1024xf32>,
               %R: memref<16x64x1024xf32>, %out: memref<1024x1024xf32>) {
   %rhs = memref.alloc() :  memref<1024x1024xf32>
 
   // Reshape from 3-d to 2-d.
   affine.for %i0 = 0 to 16 {
     affine.for %i1 = 0 to 64 {
       affine.for %k = 0 to 1024 {
         %v = affine.load %R[%i0, %i1, %k] : memref<16x64x1024xf32>
         affine.store %v, %rhs[64*%i0 + %i1, %k] : memref<1024x1024xf32>
       }
     }
   }
 
   // Matmul.
   affine.for %i = 0 to 1024 {
     affine.for %j = 0 to 1024 {
       affine.for %k = 0 to 1024 {
         %0 = affine.load %rhs[%k, %j] : memref<1024x1024xf32>
         %1 = affine.load %lhs[%i, %k] : memref<1024x1024xf32>
         %2 = arith.mulf %1, %0 : f32
         %3 = affine.load %out[%i, %j] : memref<1024x1024xf32>
         %4 = arith.addf %3, %2 : f32
         affine.store %4, %out[%i, %j] : memref<1024x1024xf32>
       }
     }
   }
   return
 }
 // MAXIMAL-NEXT: memref.alloc
 // MAXIMAL-NEXT: affine.for
 // MAXIMAL-NEXT:   affine.for
 // MAXIMAL-NEXT:     affine.for
 // MAXIMAL-NOT:      affine.for
 // MAXIMAL:      return
 
 // -----
 
 // CHECK-LABEL: func @vector_loop
 func.func @vector_loop(%a : memref<10x20xf32>, %b : memref<10x20xf32>,
                   %c : memref<10x20xf32>) {
   affine.for %j = 0 to 10 {
     affine.for %i = 0 to 5 {
       %ld0 = affine.vector_load %a[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
       affine.vector_store %ld0, %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
     }
   }
 
   affine.for %j = 0 to 10 {
     affine.for %i = 0 to 5 {
       %ld0 = affine.vector_load %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
       affine.vector_store %ld0, %c[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
     }
   }
 
   return
 }
 // CHECK:      affine.for
 // CHECK-NEXT:   affine.for
 // CHECK-NEXT:     affine.vector_load
 // CHECK-NEXT:     affine.vector_store
 // CHECK-NEXT:     affine.vector_load
 // CHECK-NEXT:     affine.vector_store
 // CHECK-NOT:  affine.for
 
 // -----
 
 // CHECK-LABEL: func @multi_outgoing_edges
 func.func @multi_outgoing_edges(%in0 : memref<32xf32>,
                       %in1 : memref<32xf32>) {
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.addf %lhs, %rhs : f32
     affine.store %add, %in0[%d] : memref<32xf32>
   }
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.subf %lhs, %rhs : f32
     affine.store %add, %in0[%d] : memref<32xf32>
   }
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.mulf %lhs, %rhs : f32
     affine.store %add, %in0[%d] : memref<32xf32>
   }
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.divf %lhs, %rhs : f32
     affine.store %add, %in0[%d] : memref<32xf32>
   }
   return
 }
 
 // CHECK:      affine.for
 // CHECK-NOT:  affine.for
 // CHECK:        arith.addf
 // CHECK-NOT:  affine.for
 // CHECK:        arith.subf
 // CHECK-NOT:  affine.for
 // CHECK:        arith.mulf
 // CHECK-NOT:  affine.for
 // CHECK:        arith.divf
 
 // -----
 
 // Test fusion when dynamically shaped memrefs are used with constant trip count loops.
 
 // CHECK-LABEL: func @calc
 func.func @calc(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>, %len: index) {
   %c1 = arith.constant 1 : index
   %1 = memref.alloc(%len) : memref<?xf32>
   affine.for %arg4 = 1 to 10 {
     %7 = affine.load %arg0[%arg4] : memref<?xf32>
     %8 = affine.load %arg1[%arg4] : memref<?xf32>
     %9 = arith.addf %7, %8 : f32
     affine.store %9, %1[%arg4] : memref<?xf32>
   }
   affine.for %arg4 = 1 to 10 {
     %7 = affine.load %1[%arg4] : memref<?xf32>
     %8 = affine.load %arg1[%arg4] : memref<?xf32>
     %9 = arith.mulf %7, %8 : f32
     affine.store %9, %arg2[%arg4] : memref<?xf32>
   }
   return
 }
 // CHECK:       memref.alloc() : memref<1xf32>
 // CHECK:       affine.for %arg{{.*}} = 1 to 10 {
 // CHECK-NEXT:    affine.load %arg{{.*}}
 // CHECK-NEXT:    affine.load %arg{{.*}}
 // CHECK-NEXT:    arith.addf
 // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
 // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
 // CHECK-NEXT:    affine.load %arg{{.*}}[%arg{{.*}}] : memref<?xf32>
 // CHECK-NEXT:    arith.mulf
 // CHECK-NEXT:    affine.store %{{.*}}, %arg{{.*}}[%arg{{.*}}] : memref<?xf32>
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
 
 // -----
 
 // CHECK-LABEL: func @should_not_fuse_since_non_affine_users
 func.func @should_not_fuse_since_non_affine_users(%in0 : memref<32xf32>,
                       %in1 : memref<32xf32>) {
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.addf %lhs, %rhs : f32
     affine.store %add, %in0[%d] : memref<32xf32>
   }
   affine.for %d = 0 to 32 {
     %lhs = memref.load %in0[%d] : memref<32xf32>
     %rhs = memref.load %in1[%d] : memref<32xf32>
     %add = arith.subf %lhs, %rhs : f32
     memref.store %add, %in0[%d] : memref<32xf32>
   }
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.mulf %lhs, %rhs : f32
     affine.store %add, %in0[%d] : memref<32xf32>
   }
   return
 }
 
 // CHECK:  affine.for
 // CHECK:    arith.addf
 // CHECK:  affine.for
 // CHECK:    arith.subf
 // CHECK:  affine.for
 // CHECK:    arith.mulf
 
 // -----
 
 // CHECK-LABEL: func @should_not_fuse_since_top_level_non_affine_users
 func.func @should_not_fuse_since_top_level_non_affine_users(%in0 : memref<32xf32>,
                       %in1 : memref<32xf32>) {
   %sum = memref.alloc() : memref<f32>
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.addf %lhs, %rhs : f32
     memref.store %add, %sum[] : memref<f32>
     affine.store %add, %in0[%d] : memref<32xf32>
   }
   %load_sum = memref.load %sum[] : memref<f32>
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.mulf %lhs, %rhs : f32
     %sub = arith.subf %add, %load_sum: f32
     affine.store %sub, %in0[%d] : memref<32xf32>
   }
   memref.dealloc %sum : memref<f32>
   return
 }
 
 // CHECK:  affine.for
 // CHECK:    arith.addf
 // CHECK:  affine.for
 // CHECK:    arith.mulf
 // CHECK:    arith.subf
 
 // -----
 
 // CHECK-LABEL: func @should_not_fuse_since_top_level_non_affine_mem_write_users
 func.func @should_not_fuse_since_top_level_non_affine_mem_write_users(
     %in0 : memref<32xf32>, %in1 : memref<32xf32>) {
   %c0 = arith.constant 0 : index
   %cst_0 = arith.constant 0.000000e+00 : f32
 
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.addf %lhs, %rhs : f32
     affine.store %add, %in0[%d] : memref<32xf32>
   }
   memref.store %cst_0, %in0[%c0] : memref<32xf32>
   affine.for %d = 0 to 32 {
     %lhs = affine.load %in0[%d] : memref<32xf32>
     %rhs = affine.load %in1[%d] : memref<32xf32>
     %add = arith.addf %lhs, %rhs: f32
     affine.store %add, %in0[%d] : memref<32xf32>
   }
   return
 }
 
 // CHECK:  affine.for
 // CHECK:    arith.addf
 // CHECK:  affine.for
 // CHECK:    arith.addf
 
 // -----
 
 // MAXIMAL-LABEL: func @fuse_minor_affine_map
 func.func @fuse_minor_affine_map(%in: memref<128xf32>, %out: memref<20x512xf32>) {
   %tmp = memref.alloc() : memref<128xf32>
 
   affine.for %arg4 = 0 to 128 {
     %ld = affine.load %in[%arg4] : memref<128xf32>
     affine.store %ld, %tmp[%arg4] : memref<128xf32>
   }
 
   affine.for %arg3 = 0 to 20 {
     affine.for %arg4 = 0 to 512 {
       %ld = affine.load %tmp[%arg4 mod 128] : memref<128xf32>
       affine.store %ld, %out[%arg3, %arg4] : memref<20x512xf32>
     }
   }
 
   return
 }
 
 // TODO: The size of the private memref is not properly computed in the presence
 // of the 'mod' operation. It should be memref<1xf32> instead of
 // memref<128xf32>: https://bugs.llvm.org/show_bug.cgi?id=46973
 // MAXIMAL:       memref.alloc() : memref<128xf32>
 // MAXIMAL:       affine.for
 // MAXIMAL-NEXT:    affine.for
 // MAXIMAL-NOT:   affine.for
 // MAXIMAL:       return
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_multi_store_producer_and_privatize_memfefs
 func.func @should_fuse_multi_store_producer_and_privatize_memfefs() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %c = memref.alloc() : memref<10xf32>
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %arg0 = 0 to 10 {
     affine.store %cst, %a[%arg0] : memref<10xf32>
     affine.store %cst, %b[%arg0] : memref<10xf32>
     affine.store %cst, %c[%arg0] : memref<10xf32>
     %0 = affine.load %c[%arg0] : memref<10xf32>
   }
 
   affine.for %arg0 = 0 to 10 {
     %0 = affine.load %a[%arg0] : memref<10xf32>
   }
 
   affine.for %arg0 = 0 to 10 {
     %0 = affine.load %b[%arg0] : memref<10xf32>
   }
 
   // All the memrefs should be privatized except '%c', which is not involved in
   // the producer-consumer fusion.
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
 
   return
 }
 
 
 func.func @should_fuse_multi_store_producer_with_escaping_memrefs_and_remove_src(
     %a : memref<10xf32>, %b : memref<10xf32>) {
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %a[%i0] : memref<10xf32>
     affine.store %cst, %b[%i0] : memref<10xf32>
   }
 
   affine.for %i1 = 0 to 10 {
     %0 = affine.load %a[%i1] : memref<10xf32>
   }
 
   affine.for %i2 = 0 to 10 {
     %0 = affine.load %b[%i2] : memref<10xf32>
   }
 
   // Producer loop '%i0' should be removed after fusion since fusion is maximal.
   // No memref should be privatized since they escape the function, and the
   // producer is removed after fusion.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NOT:   affine.for
 
   return
 }
 
 // -----
 
 func.func @should_fuse_multi_store_producer_with_escaping_memrefs_and_preserve_src(
     %a : memref<10xf32>, %b : memref<10xf32>) {
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %a[%i0] : memref<10xf32>
     affine.store %cst, %b[%i0] : memref<10xf32>
   }
 
   affine.for %i1 = 0 to 5 {
     %0 = affine.load %a[%i1] : memref<10xf32>
   }
 
   affine.for %i2 = 0 to 10 {
     %0 = affine.load %b[%i2] : memref<10xf32>
   }
 
   // Loops '%i0' and '%i2' should be fused first and '%i0' should be removed
   // since fusion is maximal. Then the fused loop and '%i1' should be fused
   // and the fused loop shouldn't be removed since fusion is not maximal.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK:       affine.for %{{.*}} = 0 to 5 {
   // CHECK-NEXT:    affine.store %{{.*}} : memref<1xf32>
   // CHECK-NEXT:    affine.store %{{.*}} : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}} : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}} : memref<1xf32>
   // CHECK-NEXT:  }
   // CHECK-NOT:   affine.for
 
   return
 }
 
 
 func.func @should_not_fuse_due_to_dealloc(%arg0: memref<16xf32>){
   %A = memref.alloc() : memref<16xf32>
   %C = memref.alloc() : memref<16xf32>
   %cst_1 = arith.constant 1.000000e+00 : f32
   affine.for %arg1 = 0 to 16 {
     %a = affine.load %arg0[%arg1] : memref<16xf32>
     affine.store %a, %A[%arg1] : memref<16xf32>
     affine.store %a, %C[%arg1] : memref<16xf32>
   }
   memref.dealloc %C : memref<16xf32>
   %B = memref.alloc() : memref<16xf32>
   affine.for %arg1 = 0 to 16 {
     %a = affine.load %A[%arg1] : memref<16xf32>
     %b = arith.addf %cst_1, %a : f32
     affine.store %b, %B[%arg1] : memref<16xf32>
   }
   memref.dealloc %A : memref<16xf32>
   return
 }
 // CHECK-LABEL: func @should_not_fuse_due_to_dealloc
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.load
 // CHECK-NEXT:      affine.store
 // CHECK-NEXT:      affine.store
 // CHECK:         memref.dealloc
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.load
 // CHECK-NEXT:      arith.addf
 // CHECK-NEXT:      affine.store
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_defining_node_has_no_dependence_from_source_node
 func.func @should_fuse_defining_node_has_no_dependence_from_source_node(
     %a : memref<10xf32>, %b : memref<f32>) -> () {
   affine.for %i0 = 0 to 10 {
     %0 = affine.load %b[] : memref<f32>
     affine.store %0, %a[%i0] : memref<10xf32>
   }
   %0 = affine.load %b[] : memref<f32>
   affine.for %i1 = 0 to 10 {
     %1 = affine.load %a[%i1] : memref<10xf32>
     %2 = arith.divf %0, %1 : f32
   }
 
   // Loops '%i0' and '%i1' should be fused even though there is a defining node
   // between the loops. It is because the node has no dependence from '%i0'.
   // CHECK:       affine.load %{{.*}}[] : memref<f32>
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[] : memref<f32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    arith.divf
   // CHECK-NEXT:  }
   // CHECK-NOT:   affine.for
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_not_fuse_defining_node_has_dependence_from_source_loop
 func.func @should_not_fuse_defining_node_has_dependence_from_source_loop(
     %a : memref<10xf32>, %b : memref<f32>) -> () {
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %b[] : memref<f32>
     affine.store %cst, %a[%i0] : memref<10xf32>
   }
   %0 = affine.load %b[] : memref<f32>
   affine.for %i1 = 0 to 10 {
     %1 = affine.load %a[%i1] : memref<10xf32>
     %2 = arith.divf %0, %1 : f32
   }
 
   // Loops '%i0' and '%i1' should not be fused because the defining node of '%0'
   // used in '%i1' has dependence from loop '%i0'.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[] : memref<f32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.load %{{.*}}[] : memref<f32>
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    arith.divf
   // CHECK-NEXT:  }
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_not_fuse_defining_node_has_transitive_dependence_from_source_loop
 func.func @should_not_fuse_defining_node_has_transitive_dependence_from_source_loop(
     %a : memref<10xf32>, %b : memref<10xf32>, %c : memref<f32>) -> () {
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %a[%i0] : memref<10xf32>
     affine.store %cst, %b[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %1 = affine.load %b[%i1] : memref<10xf32>
     affine.store %1, %c[] : memref<f32>
   }
   %0 = affine.load %c[] : memref<f32>
   affine.for %i2 = 0 to 10 {
     %1 = affine.load %a[%i2] : memref<10xf32>
     %2 = arith.divf %0, %1 : f32
   }
 
   // When loops '%i0' and '%i2' are evaluated first, they should not be
   // fused. The defining node of '%0' in loop '%i2' has transitive dependence
   // from loop '%i0'. After that, loops '%i0' and '%i1' are evaluated, and they
   // will be fused as usual.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[] : memref<f32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.load %{{.*}}[] : memref<f32>
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    arith.divf
   // CHECK-NEXT:  }
   // CHECK-NOT:   affine.for
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_not_fuse_dest_loop_nest_return_value
 func.func @should_not_fuse_dest_loop_nest_return_value(
     %a : memref<10xf32>) -> () {
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %a[%i0] : memref<10xf32>
   }
   %b = affine.for %i1 = 0 to 10 step 2 iter_args(%b_iter = %cst) -> f32 {
     %load_a = affine.load %a[%i1] : memref<10xf32>
     affine.yield %load_a: f32
   }
 
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK:       affine.for %{{.*}} = 0 to 10 step 2 iter_args(%{{.*}} = %{{.*}}) -> (f32) {
   // CHECK-NEXT:    affine.load
   // CHECK-NEXT:    affine.yield
   // CHECK-NEXT:  }
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_not_fuse_src_loop_nest_return_value
 func.func @should_not_fuse_src_loop_nest_return_value(
     %a : memref<10xf32>) -> () {
   %cst = arith.constant 1.000000e+00 : f32
   %b = affine.for %i = 0 to 10 step 2 iter_args(%b_iter = %cst) -> f32 {
     %c = arith.addf %b_iter, %b_iter : f32
     affine.store %c, %a[%i] : memref<10xf32>
     affine.yield %c: f32
   }
   affine.for %i1 = 0 to 10 {
     %1 = affine.load %a[%i1] : memref<10xf32>
   }
 
   // CHECK:       %{{.*}} = affine.for %{{.*}} = 0 to 10 step 2 iter_args(%{{.*}} = %{{.*}}) -> (f32) {
   // CHECK-NEXT:    %{{.*}} = arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.yield %{{.*}} : f32
   // CHECK-NEXT:  }
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
 
   return
 }
 
 // -----
 
 func.func private @some_function(memref<16xf32>)
 func.func @call_op_prevents_fusion(%arg0: memref<16xf32>){
   %A = memref.alloc() : memref<16xf32>
   %cst_1 = arith.constant 1.000000e+00 : f32
   affine.for %arg1 = 0 to 16 {
     %a = affine.load %arg0[%arg1] : memref<16xf32>
     affine.store %a, %A[%arg1] : memref<16xf32>
   }
   call @some_function(%A) : (memref<16xf32>) -> ()
   %B = memref.alloc() : memref<16xf32>
   affine.for %arg1 = 0 to 16 {
     %a = affine.load %A[%arg1] : memref<16xf32>
     %b = arith.addf %cst_1, %a : f32
     affine.store %b, %B[%arg1] : memref<16xf32>
   }
   return
 }
 // CHECK-LABEL: func @call_op_prevents_fusion
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.load
 // CHECK-NEXT:      affine.store
 // CHECK:         call
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.load
 // CHECK-NEXT:      arith.addf
 // CHECK-NEXT:      affine.store
 
 // -----
 
 func.func private @some_function()
 func.func @call_op_does_not_prevent_fusion(%arg0: memref<16xf32>){
   %A = memref.alloc() : memref<16xf32>
   %cst_1 = arith.constant 1.000000e+00 : f32
   affine.for %arg1 = 0 to 16 {
     %a = affine.load %arg0[%arg1] : memref<16xf32>
     affine.store %a, %A[%arg1] : memref<16xf32>
   }
   call @some_function() : () -> ()
   %B = memref.alloc() : memref<16xf32>
   affine.for %arg1 = 0 to 16 {
     %a = affine.load %A[%arg1] : memref<16xf32>
     %b = arith.addf %cst_1, %a : f32
     affine.store %b, %B[%arg1] : memref<16xf32>
   }
   return
 }
 // CHECK-LABEL: func @call_op_does_not_prevent_fusion
 // CHECK:         affine.for
 // CHECK-NOT:     affine.for
 
 // -----
 
 // Test for source that writes to an escaping memref and has two consumers.
 // Fusion should create private memrefs in place of `%arg0` since the source is
 // not to be removed after fusion and the destinations do not write to `%arg0`.
 // This should enable both the consumers to benefit from fusion, which would not
 // be possible if private memrefs were not created.
 func.func @should_fuse_with_both_consumers_separately(%arg0: memref<10xf32>) {
   %cf7 = arith.constant 7.0 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %arg0[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 7 {
     %v0 = affine.load %arg0[%i1] : memref<10xf32>
   }
   affine.for %i1 = 5 to 9 {
     %v0 = affine.load %arg0[%i1] : memref<10xf32>
   }
   return
 }
 
 // CHECK-LABEL: func @should_fuse_with_both_consumers_separately
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.store
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.store
 // CHECK-NEXT:      affine.load
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.store
 // CHECK-NEXT:      affine.load
 
 // -----
 
 // Fusion is avoided when the slice computed is invalid. Comments below describe
 // incorrect backward slice computation. Similar logic applies for forward slice
 // as well.
 func.func @no_fusion_cannot_compute_valid_slice() {
   %A = memref.alloc() : memref<5xf32>
   %B = memref.alloc() : memref<6xf32>
   %C = memref.alloc() : memref<5xf32>
   %cst = arith.constant 0. : f32
 
   affine.for %arg0 = 0 to 5 {
     %a = affine.load %A[%arg0] : memref<5xf32>
     affine.store %a, %B[%arg0 + 1] : memref<6xf32>
   }
 
   affine.for %arg0 = 0 to 5 {
     // Backward slice computed will be:
     // slice ( src loop: 0, dst loop: 1, depth: 1 : insert point: (1, 0)
     // loop bounds: [(d0) -> (d0 - 1), (d0) -> (d0)] )
 
     // Resulting fusion would be as below. It is easy to note the out-of-bounds
     // access by 'affine.load'.
 
     // #map0 = affine_map<(d0) -> (d0 - 1)>
     // #map1 = affine_map<(d0) -> (d0)>
     // affine.for %arg1 = #map0(%arg0) to #map1(%arg0) {
     //   %5 = affine.load %1[%arg1] : memref<5xf32>
     //   ...
     //   ...
     // }
 
     %a = affine.load %B[%arg0] : memref<6xf32>
     %b = arith.mulf %a, %cst : f32
     affine.store %b, %C[%arg0] : memref<5xf32>
   }
   return
 }
 // CHECK-LABEL: func @no_fusion_cannot_compute_valid_slice
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.load
 // CHECK-NEXT:      affine.store
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.load
 // CHECK-NEXT:      arith.mulf
 // CHECK-NEXT:      affine.store
 
 // MAXIMAL-LABEL:   func @reduce_add_f32_f32(
 func.func @reduce_add_f32_f32(%arg0: memref<64x64xf32, 1>, %arg1: memref<1x64xf32, 1>, %arg2: memref<1x64xf32, 1>) {
   %cst_0 = arith.constant 0.000000e+00 : f32
   %cst_1 = arith.constant 1.000000e+00 : f32
   %0 = memref.alloca() : memref<f32, 1>
   %1 = memref.alloca() : memref<f32, 1>
   affine.for %arg3 = 0 to 1 {
     affine.for %arg4 = 0 to 64 {
       %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_0) -> f32 {
         %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
         %5 = arith.addf %prevAccum, %4 : f32
         affine.yield %5 : f32
       }
       %accum_dbl = arith.addf %accum, %accum : f32
       affine.store %accum_dbl, %arg1[%arg3, %arg4] : memref<1x64xf32, 1>
     }
   }
   affine.for %arg3 = 0 to 1 {
     affine.for %arg4 = 0 to 64 {
       %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_1) -> f32 {
         %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
         %5 = arith.mulf %prevAccum, %4 : f32
         affine.yield %5 : f32
       }
       %accum_sqr = arith.mulf %accum, %accum : f32
       affine.store %accum_sqr, %arg2[%arg3, %arg4] : memref<1x64xf32, 1>
     }
   }
   return
 }
 // The two loops here get maximally sibling-fused at the innermost
 // insertion point. Test checks  if the innermost reduction loop of the fused loop
 // gets promoted into its outerloop.
 // MAXIMAL-SAME:                             %[[arg_0:.*]]: memref<64x64xf32, 1>,
 // MAXIMAL-SAME:                             %[[arg_1:.*]]: memref<1x64xf32, 1>,
 // MAXIMAL-SAME:                             %[[arg_2:.*]]: memref<1x64xf32, 1>) {
 // MAXIMAL:             %[[cst:.*]] = arith.constant 0 : index
 // MAXIMAL-NEXT:        %[[cst_0:.*]] = arith.constant 0.000000e+00 : f32
 // MAXIMAL-NEXT:        %[[cst_1:.*]] = arith.constant 1.000000e+00 : f32
 // MAXIMAL:             affine.for %[[idx_0:.*]] = 0 to 1 {
 // MAXIMAL-NEXT:          affine.for %[[idx_1:.*]] = 0 to 64 {
 // MAXIMAL-NEXT:            %[[results:.*]]:2 = affine.for %[[idx_2:.*]] = 0 to 64 iter_args(%[[iter_0:.*]] = %[[cst_1]], %[[iter_1:.*]] = %[[cst_0]]) -> (f32, f32) {
 // MAXIMAL-NEXT:              %[[val_0:.*]] = affine.load %[[arg_0]][%[[idx_2]], %[[idx_1]]] : memref<64x64xf32, 1>
 // MAXIMAL-NEXT:              %[[reduc_0:.*]] = arith.addf %[[iter_1]], %[[val_0]] : f32
 // MAXIMAL-NEXT:              %[[val_1:.*]] = affine.load %[[arg_0]][%[[idx_2]], %[[idx_1]]] : memref<64x64xf32, 1>
 // MAXIMAL-NEXT:              %[[reduc_1:.*]] = arith.mulf %[[iter_0]], %[[val_1]] : f32
 // MAXIMAL-NEXT:              affine.yield %[[reduc_1]], %[[reduc_0]] : f32, f32
 // MAXIMAL-NEXT:            }
 // MAXIMAL-NEXT:            %[[reduc_0_dbl:.*]] = arith.addf %[[results:.*]]#1, %[[results]]#1 : f32
 // MAXIMAL-NEXT:            affine.store %[[reduc_0_dbl]], %[[arg_1]][%[[cst]], %[[idx_1]]] : memref<1x64xf32, 1>
 // MAXIMAL-NEXT:            %[[reduc_1_sqr:.*]] = arith.mulf %[[results]]#0, %[[results]]#0 : f32
 // MAXIMAL-NEXT:            affine.store %[[reduc_1_sqr]], %[[arg_2]][%[[idx_0]], %[[idx_1]]] : memref<1x64xf32, 1>
 // MAXIMAL-NEXT:          }
 // MAXIMAL-NEXT:        }
 // MAXIMAL-NEXT:        return
 // MAXIMAL-NEXT:      }
 
 // -----
 
 // CHECK-LABEL:   func @reduce_add_non_innermost
 func.func @reduce_add_non_innermost(%arg0: memref<64x64xf32, 1>, %arg1: memref<1x64xf32, 1>, %arg2: memref<1x64xf32, 1>) {
   %cst = arith.constant 0.000000e+00 : f32
   %cst_0 = arith.constant 1.000000e+00 : f32
   %0 = memref.alloca() : memref<f32, 1>
   %1 = memref.alloca() : memref<f32, 1>
   affine.for %arg3 = 0 to 1 {
     affine.for %arg4 = 0 to 64 {
       %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst) -> f32 {
         %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
         %5 = arith.addf %prevAccum, %4 : f32
         affine.yield %5 : f32
       }
       %accum_dbl = arith.addf %accum, %accum : f32
       affine.store %accum_dbl, %arg1[%arg3, %arg4] : memref<1x64xf32, 1>
     }
   }
   affine.for %arg3 = 0 to 1 {
     affine.for %arg4 = 0 to 64 {
       %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_0) -> f32 {
         %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
         %5 = arith.mulf %prevAccum, %4 : f32
         affine.yield %5 : f32
       }
       %accum_sqr = arith.mulf %accum, %accum : f32
       affine.store %accum_sqr, %arg2[%arg3, %arg4] : memref<1x64xf32, 1>
     }
   }
   return
 }
 // Test checks the loop structure is preserved after sibling fusion.
 // CHECK:         affine.for
 // CHECK-NEXT:      affine.for
 // CHECK-NEXT:        affine.for
 // CHECK             affine.for
 
 
 
 // -----
 
 // CHECK-LABEL: func @fuse_large_number_of_loops
 func.func @fuse_large_number_of_loops(%arg0: memref<20x10xf32, 1>, %arg1: memref<20x10xf32, 1>, %arg2: memref<20x10xf32, 1>, %arg3: memref<20x10xf32, 1>, %arg4: memref<20x10xf32, 1>, %arg5: memref<f32, 1>, %arg6: memref<f32, 1>, %arg7: memref<f32, 1>, %arg8: memref<f32, 1>, %arg9: memref<20x10xf32, 1>, %arg10: memref<20x10xf32, 1>, %arg11: memref<20x10xf32, 1>, %arg12: memref<20x10xf32, 1>) {
   %cst = arith.constant 1.000000e+00 : f32
   %0 = memref.alloc() : memref<f32, 1>
   affine.store %cst, %0[] : memref<f32, 1>
   %1 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg6[] : memref<f32, 1>
       affine.store %21, %1[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %2 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %1[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %arg3[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.mulf %22, %21 : f32
       affine.store %23, %2[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %3 = memref.alloc() : memref<f32, 1>
   %4 = affine.load %arg6[] : memref<f32, 1>
   %5 = affine.load %0[] : memref<f32, 1>
   %6 = arith.subf %5, %4 : f32
   affine.store %6, %3[] : memref<f32, 1>
   %7 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %3[] : memref<f32, 1>
       affine.store %21, %7[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %8 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg1[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %7[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.mulf %22, %21 : f32
       affine.store %23, %8[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %9 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg1[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %8[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.mulf %22, %21 : f32
       affine.store %23, %9[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %9[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %2[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.addf %22, %21 : f32
       affine.store %23, %arg11[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %10 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %1[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %arg2[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.mulf %22, %21 : f32
       affine.store %23, %10[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %8[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %10[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.addf %22, %21 : f32
       affine.store %23, %arg10[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %11 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg10[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %arg10[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.mulf %22, %21 : f32
       affine.store %23, %11[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %12 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %11[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %arg11[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.subf %22, %21 : f32
       affine.store %23, %12[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %13 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg7[] : memref<f32, 1>
       affine.store %21, %13[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %14 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg4[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %13[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.mulf %22, %21 : f32
       affine.store %23, %14[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %15 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg8[] : memref<f32, 1>
       affine.store %21, %15[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %16 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %15[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %12[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.addf %22, %21 : f32
       affine.store %23, %16[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %17 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %16[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = math.sqrt %21 : f32
       affine.store %22, %17[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %18 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg5[] : memref<f32, 1>
       affine.store %21, %18[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %19 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg1[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %18[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.mulf %22, %21 : f32
       affine.store %23, %19[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   %20 = memref.alloc() : memref<20x10xf32, 1>
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %17[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %19[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.divf %22, %21 : f32
       affine.store %23, %20[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %20[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %14[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.addf %22, %21 : f32
       affine.store %23, %arg12[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   affine.for %arg13 = 0 to 20 {
     affine.for %arg14 = 0 to 10 {
       %21 = affine.load %arg12[%arg13, %arg14] : memref<20x10xf32, 1>
       %22 = affine.load %arg0[%arg13, %arg14] : memref<20x10xf32, 1>
       %23 = arith.subf %22, %21 : f32
       affine.store %23, %arg9[%arg13, %arg14] : memref<20x10xf32, 1>
     }
   }
   return
 }
 // CHECK:         affine.for
 // CHECK:         affine.for
 // CHECK-NOT:     affine.for
 
 // CHECK-LABEL: func @alias_escaping_memref
 func.func @alias_escaping_memref(%a : memref<2x5xf32>) {
   %cst = arith.constant 0.000000e+00 : f32
   %alias = memref.reinterpret_cast %a to offset: [0], sizes: [10], strides: [1] : memref<2x5xf32> to memref<10xf32>
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %alias[%i0] : memref<10xf32>
   }
 
   affine.for %i1 = 0 to 10 {
     %0 = affine.load %alias[%i1] : memref<10xf32>
   }
   // Fusion happens, but memref isn't privatized since %alias is an alias of a
   // function argument.
   // CHECK:       memref.reinterpret_cast
   // CHECK-NEXT:  affine.for
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NOT:   affine.for
 
   return
 }
 
 // CHECK-LABEL: func @unknown_memref_def_op
 func.func @unknown_memref_def_op() {
   %cst = arith.constant 0.000000e+00 : f32
   %may_alias = call @bar() : () -> memref<10xf32>
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %may_alias[%i0] : memref<10xf32>
   }
 
   affine.for %i1 = 0 to 10 {
     %0 = affine.load %may_alias[%i1] : memref<10xf32>
   }
   // Fusion happens, but memref isn't privatized since %may_alias's origin is
   // unknown.
   // CHECK:       call
   // CHECK-NEXT:  affine.for
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NOT:   affine.for
 
   return
 }
 func.func private @bar() -> memref<10xf32>
 
 
 // Add further tests in mlir/test/Transforms/loop-fusion-4.mlir

mlir/test/Dialect/Affine/loop-fusion-4.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer}))' -split-input-file | FileCheck %s --check-prefix=PRODUCER-CONSUMER
 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{fusion-maximal mode=sibling}))' -split-input-file | FileCheck %s --check-prefix=SIBLING-MAXIMAL
 
 // Part I of fusion tests in  mlir/test/Transforms/loop-fusion.mlir.
 // Part II of fusion tests in mlir/test/Transforms/loop-fusion-2.mlir
 // Part III of fusion tests in mlir/test/Transforms/loop-fusion-3.mlir
 
 // Expects fusion of producer into consumer at depth 4 and subsequent removal of
 // source loop.
 // PRODUCER-CONSUMER-LABEL: func @unflatten4d
 func.func @unflatten4d(%arg1: memref<7x8x9x10xf32>) {
   %m = memref.alloc() : memref<5040xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 7 {
     affine.for %i1 = 0 to 8 {
       affine.for %i2 = 0 to 9 {
         affine.for %i3 = 0 to 10 {
           affine.store %cf7, %m[720 * %i0 + 90 * %i1 + 10 * %i2 + %i3] : memref<5040xf32>
         }
       }
     }
   }
   affine.for %i0 = 0 to 7 {
     affine.for %i1 = 0 to 8 {
       affine.for %i2 = 0 to 9 {
         affine.for %i3 = 0 to 10 {
           %v0 = affine.load %m[720 * %i0 + 90 * %i1 + 10 * %i2 + %i3] : memref<5040xf32>
           affine.store %v0, %arg1[%i0, %i1, %i2, %i3] : memref<7x8x9x10xf32>
         }
       }
     }
   }
   return
 }
 
 // PRODUCER-CONSUMER:        affine.for
 // PRODUCER-CONSUMER-NEXT:     affine.for
 // PRODUCER-CONSUMER-NEXT:       affine.for
 // PRODUCER-CONSUMER-NEXT:         affine.for
 // PRODUCER-CONSUMER-NOT:    affine.for
 // PRODUCER-CONSUMER: return
 
 // -----
 
 // Expects fusion of producer into consumer at depth 2 and subsequent removal of
 // source loop.
 // PRODUCER-CONSUMER-LABEL: func @unflatten2d_with_transpose
 func.func @unflatten2d_with_transpose(%arg1: memref<8x7xf32>) {
   %m = memref.alloc() : memref<56xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 7 {
     affine.for %i1 = 0 to 8 {
       affine.store %cf7, %m[8 * %i0 + %i1] : memref<56xf32>
     }
   }
   affine.for %i0 = 0 to 8 {
     affine.for %i1 = 0 to 7 {
       %v0 = affine.load %m[%i0 + 8 * %i1] : memref<56xf32>
       affine.store %v0, %arg1[%i0, %i1] : memref<8x7xf32>
     }
   }
   return
 }
 
 // PRODUCER-CONSUMER:        affine.for
 // PRODUCER-CONSUMER-NEXT:     affine.for
 // PRODUCER-CONSUMER-NOT:    affine.for
 // PRODUCER-CONSUMER: return
 
 // -----
 
 // Expects fusion of producer into consumer at depth 1 and source loop to not
 // be removed due to difference in loop steps.
 // PRODUCER-CONSUMER-LABEL: func @check_src_dst_step
 func.func @check_src_dst_step(%m : memref<100xf32>,
                          %src: memref<100xf32>,
                          %out: memref<100xf32>) {
   affine.for %i0 = 0 to 100 {
     %r1 = affine.load %src[%i0]: memref<100xf32>
     affine.store %r1, %m[%i0] : memref<100xf32>
   }
   affine.for %i2 = 0 to 100 step 2 {
     %r2 = affine.load %m[%i2] : memref<100xf32>
     affine.store %r2, %out[%i2] : memref<100xf32>
   }
   return
 }
 
 // Check if the fusion did take place as well as that the source loop was
 // not removed. To check if fusion took place, the read instruction from the
 // original source loop is checked to be in the fused loop.
 //
 // PRODUCER-CONSUMER:        affine.for %[[idx_0:.*]] = 0 to 100 {
 // PRODUCER-CONSUMER-NEXT:     %[[result_0:.*]] = affine.load %[[arr1:.*]][%[[idx_0]]] : memref<100xf32>
 // PRODUCER-CONSUMER-NEXT:     affine.store %[[result_0]], %{{.*}}[%[[idx_0]]] : memref<100xf32>
 // PRODUCER-CONSUMER-NEXT:   }
 // PRODUCER-CONSUMER:        affine.for %[[idx_1:.*]] = 0 to 100 step 2 {
 // PRODUCER-CONSUMER:          affine.load %[[arr1]][%[[idx_1]]] : memref<100xf32>
 // PRODUCER-CONSUMER:        }
 // PRODUCER-CONSUMER:        return
 
 // -----
 
 // SIBLING-MAXIMAL-LABEL:   func @reduce_add_non_maximal_f32_f32(
 func.func @reduce_add_non_maximal_f32_f32(%arg0: memref<64x64xf32, 1>, %arg1 : memref<1x64xf32, 1>, %arg2 : memref<1x64xf32, 1>) {
     %cst_0 = arith.constant 0.000000e+00 : f32
     %cst_1 = arith.constant 1.000000e+00 : f32
     affine.for %arg3 = 0 to 1 {
       affine.for %arg4 = 0 to 64 {
         %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_0) -> f32 {
           %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
           %5 = arith.addf %prevAccum, %4 : f32
           affine.yield %5 : f32
         }
         %accum_dbl = arith.addf %accum, %accum : f32
         affine.store %accum_dbl, %arg1[%arg3, %arg4] : memref<1x64xf32, 1>
       }
     }
     affine.for %arg3 = 0 to 1 {
       affine.for %arg4 = 0 to 64 {
         // Following loop  trip count does not match the corresponding source trip count.
         %accum = affine.for %arg5 = 0 to 32 iter_args (%prevAccum = %cst_1) -> f32 {
           %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
           %5 = arith.mulf %prevAccum, %4 : f32
           affine.yield %5 : f32
         }
         %accum_sqr = arith.mulf %accum, %accum : f32
         affine.store %accum_sqr, %arg2[%arg3, %arg4] : memref<1x64xf32, 1>
       }
     }
     return
 }
 // Test checks the loop structure is preserved after sibling fusion
 // since the destination loop and source loop trip counts do not
 // match.
 // SIBLING-MAXIMAL:        %[[cst_0:.*]] = arith.constant 0.000000e+00 : f32
 // SIBLING-MAXIMAL-NEXT:        %[[cst_1:.*]] = arith.constant 1.000000e+00 : f32
 // SIBLING-MAXIMAL-NEXT:           affine.for %[[idx_0:.*]]= 0 to 1 {
 // SIBLING-MAXIMAL-NEXT:             affine.for %[[idx_1:.*]] = 0 to 64 {
 // SIBLING-MAXIMAL-NEXT:               %[[result_1:.*]] = affine.for %[[idx_2:.*]] = 0 to 32 iter_args(%[[iter_0:.*]] = %[[cst_1]]) -> (f32) {
 // SIBLING-MAXIMAL-NEXT:                 %[[result_0:.*]] = affine.for %[[idx_3:.*]] = 0 to 64 iter_args(%[[iter_1:.*]] = %[[cst_0]]) -> (f32) {
 
 // -----
 
 // SIBLING-MAXIMAL-LABEL: func @sibling_load_only
 func.func @sibling_load_only(%arg0: memref<10xf32>) {
   affine.for %arg1 = 0 to 10 {
     %0 = affine.load %arg0[%arg1] : memref<10xf32>
   }
   affine.for %arg1 = 0 to 10 {
     %0 = affine.load %arg0[%arg1] : memref<10xf32>
   }
   // SIBLING-MAXIMAL-NEXT: affine.for
   // SIBLING-MAXIMAL-NEXT:   affine.load
   // SIBLING-MAXIMAL-NEXT:   affine.load
   return
 }
 
 // -----
 
 // PRODUCER-CONSUMER-LABEL: func @fusion_for_multiple_blocks() {
 func.func @fusion_for_multiple_blocks() {
 ^bb0:
   %m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // PRODUCER-CONSUMER:      affine.for %{{.*}} = 0 to 10 {
   // PRODUCER-CONSUMER-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // PRODUCER-CONSUMER-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // PRODUCER-CONSUMER-NEXT: }
   cf.br ^bb1
 ^bb1:
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // PRODUCER-CONSUMER:      affine.for %{{.*}} = 0 to 10 {
   // PRODUCER-CONSUMER-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // PRODUCER-CONSUMER-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // PRODUCER-CONSUMER-NEXT: }
   return
 }
 
 // -----
 
 // PRODUCER-CONSUMER-LABEL: @fuse_higher_dim_nest_into_lower_dim_nest
 func.func @fuse_higher_dim_nest_into_lower_dim_nest() {
   %A = memref.alloc() : memref<8x12x128x64xf32>
   %B = memref.alloc() : memref<8x128x12x64xf32>
   affine.for %arg205 = 0 to 8 {
     affine.for %arg206 = 0 to 128 {
       affine.for %arg207 = 0 to 12 {
         affine.for %arg208 = 0 to 64 {
           %a = affine.load %A[%arg205, %arg207, %arg206, %arg208] : memref<8x12x128x64xf32>
           affine.store %a, %B[%arg205, %arg206, %arg207, %arg208] : memref<8x128x12x64xf32>
         }
       }
     }
   }
   %C = memref.alloc() : memref<8x128x768xf16>
   affine.for %arg205 = 0 to 8 {
     affine.for %arg206 = 0 to 128 {
       affine.for %arg207 = 0 to 768 {
         %b = affine.load %B[%arg205, %arg206, %arg207 floordiv 64, %arg207 mod 64] : memref<8x128x12x64xf32>
         %c = arith.truncf %b : f32 to f16
         affine.store %c, %C[%arg205, %arg206, %arg207] : memref<8x128x768xf16>
       }
     }
   }
 
   // Check that fusion happens into the innermost loop of the consumer.
   // PRODUCER-CONSUMER:      affine.for
   // PRODUCER-CONSUMER-NEXT:   affine.for %{{.*}} = 0 to 128
   // PRODUCER-CONSUMER-NEXT:     affine.for %{{.*}} = 0 to 768
   // PRODUCER-CONSUMER-NOT:  affine.for
   // PRODUCER-CONSUMER:      return
   return
 }

mlir/test/Dialect/Affine/loop-fusion.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' -split-input-file | FileCheck %s
 
 // Part II of fusion tests in  mlir/test/Transforms/loop-fusion=2.mlir.
 // Part III of fusion tests in mlir/test/Transforms/loop-fusion-3.mlir
 // Part IV of fusion tests in mlir/test/Transforms/loop-fusion-4.mlir
 
 // TODO: Add more tests:
 // *) Add nested fusion test cases when non-constant loop bound support is
 //    added to iteration domain in dependence check.
 // *) Add a test w/ floordiv/ceildiv/mod when supported in dependence check.
 // *) Add tests which check fused computation slice indexing and loop bounds.
 // TODO: Test clean up: move memref allocs to func args.
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_raw_dep_for_locality() {
 func.func @should_fuse_raw_dep_for_locality() {
   %m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_reduction_to_pointwise() {
 func.func @should_fuse_reduction_to_pointwise() {
   %a = memref.alloc() : memref<10x10xf32>
   %b = memref.alloc() : memref<10xf32>
   %c = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.for %i1 = 0 to 10 {
       %v0 = affine.load %b[%i0] : memref<10xf32>
       %v1 = affine.load %a[%i0, %i1] : memref<10x10xf32>
       %v3 = arith.addf %v0, %v1 : f32
       affine.store %v3, %b[%i0] : memref<10xf32>
     }
   }
   affine.for %i2 = 0 to 10 {
     %v4 = affine.load %b[%i2] : memref<10xf32>
     affine.store %v4, %c[%i2] : memref<10xf32>
   }
 
   // Should fuse in entire inner loop on %i1 from source loop nest, as %i1
   // is not used in the access function of the store/load on %b.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
   // CHECK-NEXT:      arith.addf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-DAG: [[$MAP_SHIFT_MINUS_ONE_R1:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 - 1)>
 // CHECK-DAG: [[$MAP_SHIFT_D0_BY_ONE:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> (d0 + 1)>
 // CHECK-DAG: [[$MAP_SHIFT_D1_BY_ONE:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> (d1 + 1)>
 
 // CHECK-LABEL: func @should_fuse_loop_nests_with_shifts() {
 func.func @should_fuse_loop_nests_with_shifts() {
   %a = memref.alloc() : memref<10x10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 9 {
     affine.for %i1 = 0 to 9 {
       affine.store %cf7, %a[%i0 + 1, %i1 + 1] : memref<10x10xf32>
     }
   }
   affine.for %i2 = 1 to 10 {
     affine.for %i3 = 1 to 10 {
       %v0 = affine.load %a[%i2, %i3] : memref<10x10xf32>
     }
   }
 
   // Source slice affine apply sequence:
   // *) First two affine apply's map from the dst to src iteration space.
   // *) Third affine apply is access function around src store.
   // *) Fourth affine apply shifts the stores access function by '-1', because
   //    of the offset induced by reducing the memref shape from 10x10 to 9x9.
   // *) Fifth affine apply shifts the loads access function by '-1', because
   //    of the offset induced by reducing the memref shape from 10x10 to 9x9.
   // NOTE: Should create a private memref with reduced shape 9x9xf32.
   // CHECK:      affine.for %{{.*}} = 1 to 10 {
   // CHECK-NEXT:   affine.for %{{.*}} = 1 to 10 {
   // CHECK-NEXT:     %[[I:.*]] = affine.apply [[$MAP_SHIFT_MINUS_ONE_R1]](%{{.*}})
   // CHECK-NEXT:     %[[J:.*]] = affine.apply [[$MAP_SHIFT_MINUS_ONE_R1]](%{{.*}})
   // CHECK-NEXT:     affine.apply [[$MAP_SHIFT_D0_BY_ONE]](%[[I]], %[[J]])
   // CHECK-NEXT:     affine.apply [[$MAP_SHIFT_D1_BY_ONE]](%[[I]], %[[J]])
   // CHECK-NEXT:     affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
   // CHECK-NEXT:     affine.load %{{.*}}[0, 0] : memref<1x1xf32>
   // CHECK-NEXT:   }
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_loop_nest() {
 func.func @should_fuse_loop_nest() {
   %a = memref.alloc() : memref<10x10xf32>
   %b = memref.alloc() : memref<10x10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.for %i1 = 0 to 10 {
       affine.store %cf7, %a[%i0, %i1] : memref<10x10xf32>
     }
   }
   affine.for %i2 = 0 to 10 {
     affine.for %i3 = 0 to 10 {
       %v0 = affine.load %a[%i3, %i2] : memref<10x10xf32>
       affine.store %v0, %b[%i2, %i3] : memref<10x10xf32>
     }
   }
   affine.for %i4 = 0 to 10 {
     affine.for %i5 = 0 to 10 {
       %v1 = affine.load %b[%i4, %i5] : memref<10x10xf32>
     }
   }
   // Expecting private memref for '%a' first, then private memref for '%b'.
   // CHECK-DAG:  [[NEWA:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1x1xf32>
   // CHECK-DAG:  [[NEWB:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1x1xf32>
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:     affine.store %{{.*}}, [[NEWA]][0, 0] : memref<1x1xf32>
   // CHECK-NEXT:     affine.load [[NEWA]][0, 0] : memref<1x1xf32>
   // CHECK-NEXT:     affine.store %{{.*}}, [[NEWB]][0, 0] : memref<1x1xf32>
   // CHECK-NEXT:     affine.load [[NEWB]][0, 0] : memref<1x1xf32>
   // CHECK-NEXT:   }
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_across_intermediate_loop_with_no_deps() {
 func.func @should_fuse_across_intermediate_loop_with_no_deps() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %c = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     %v0 = affine.load %a[%i0] : memref<10xf32>
     affine.store %v0, %b[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.store %cf7, %c[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     %v1 = affine.load %b[%i2] : memref<10xf32>
   }
 
   // Should fuse first loop (past second loop with no dependences) into third.
   // Note that fusion creates a private memref '%2' for the fused loop nest.
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_all_loops() {
 func.func @should_fuse_all_loops() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   // Set up flow dependences from first and second loops to third.
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %a[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.store %cf7, %b[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     %v0 = affine.load %a[%i2] : memref<10xf32>
     %v1 = affine.load %b[%i2] : memref<10xf32>
   }
 
   // Should fuse first and second loops into third.
   // Expecting private memref for '%a' first, then private memref for '%b'.
   // CHECK-DAG: [[NEWA:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1xf32>
   // CHECK-DAG: [[NEWB:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1xf32>
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, [[NEWA]][0] : memref<1xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, [[NEWB]][0] : memref<1xf32>
   // CHECK-NEXT:   affine.load [[NEWA]][0] : memref<1xf32>
   // CHECK-NEXT:   affine.load [[NEWB]][0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_first_and_second_loops() {
 func.func @should_fuse_first_and_second_loops() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %c = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %a[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %a[%i1] : memref<10xf32>
     affine.store %cf7, %b[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     %v1 = affine.load %c[%i2] : memref<10xf32>
   }
 
   // Should fuse first loop into the second (last loop should not be fused).
   // Should create private memref '%2' for fused scf.
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_not_fuse_would_create_cycle() {
 func.func @should_not_fuse_would_create_cycle() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %c = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
 
   // Set up the following dependences:
   // 1) loop0 -> loop1 on memref '%{{.*}}'
   // 2) loop0 -> loop2 on memref '%{{.*}}'
   // 3) loop1 -> loop2 on memref '%{{.*}}'
   affine.for %i0 = 0 to 10 {
     %v0 = affine.load %a[%i0] : memref<10xf32>
     affine.store %cf7, %b[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.store %cf7, %a[%i1] : memref<10xf32>
     %v1 = affine.load %c[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     %v2 = affine.load %b[%i2] : memref<10xf32>
     affine.store %cf7, %c[%i2] : memref<10xf32>
   }
   // Should not fuse: fusing loop first loop into last would create a cycle.
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_producer_consumer() {
 func.func @should_fuse_producer_consumer() {
   %m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.store %cf7, %m[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     %v1 = affine.load %m[%i2] : memref<10xf32>
   }
   // Fusing loop %i0 to %i2 would violate the WAW dependence between %i0 and
   // %i1, but OK to fuse %i1 into %i2.
   // TODO: When the fusion pass is run to a fixed-point, it should
   // fuse all three of these loop nests.
   // CHECK:      memref.alloc() : memref<1xf32>
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_and_move_to_preserve_war_dep() {
 func.func @should_fuse_and_move_to_preserve_war_dep() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     %v0 = affine.load %a[%i0] : memref<10xf32>
     affine.store %v0, %b[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.store %cf7, %a[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     %v1 = affine.load %b[%i2] : memref<10xf32>
   }
   // Loops '%i1' and '%i2' have no dependences. We can fuse a slice of '%i0'
   // into '%i2' if we move the fused loop nest before '%i1', which preserves
   // the WAR dependence from load '%a' in '%i0' to the store '%a' in loop '%i1'.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_if_top_level_access() {
 func.func @should_fuse_if_top_level_access() {
   %m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
 
   %c0 = arith.constant 4 : index
   %v1 = affine.load %m[%c0] : memref<10xf32>
   // Top-level load to '%m' should prevent creating a private memref but
   // loop nests should be fused and '%i0' should be removed.
   // CHECK:      %[[m:.*]] = memref.alloc() : memref<10xf32>
   // CHECK-NOT:  memref.alloc
 
   // CHECK:      affine.for %[[i1:.*]] = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %[[m]][%[[i1]]] : memref<10xf32>
   // CHECK-NEXT:   affine.load %[[m]][%[[i1]]] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK:      affine.load %[[m]][%{{.*}}] : memref<10xf32>
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_but_not_remove_src() {
 func.func @should_fuse_but_not_remove_src() {
   %m = memref.alloc() : memref<100xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 100 {
     affine.store %cf7, %m[%i0] : memref<100xf32>
   }
   affine.for %i1 = 0 to 17 {
     %v0 = affine.load %m[%i1] : memref<100xf32>
   }
   %v1 = affine.load %m[99] : memref<100xf32>
 
   // Loop '%i0' and '%i1' should be fused but '%i0' shouldn't be removed to
   // preserve the dependence with the top-level access.
   // CHECK:      affine.for %{{.*}} = 0 to 100 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<100xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: affine.for %{{.*}} = 0 to 17 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: affine.load %{{.*}}[99] : memref<100xf32>
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_no_top_level_access() {
 func.func @should_fuse_no_top_level_access() {
   %m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 #set0 = affine_set<(d0) : (1 == 0)>
 
 // CHECK-LABEL: func @should_not_fuse_if_op_at_top_level() {
 func.func @should_not_fuse_if_op_at_top_level() {
   %m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   %c0 = arith.constant 4 : index
   affine.if #set0(%c0) {
   }
   // Top-level IfOp should prevent fusion.
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   return
 }
 
 // -----
 
 #set0 = affine_set<(d0) : (1 == 0)>
 
 // CHECK-LABEL: func @should_not_fuse_if_op_in_loop_nest() {
 func.func @should_not_fuse_if_op_in_loop_nest() {
   %m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
   %c4 = arith.constant 4 : index
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.if #set0(%c4) {
     }
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
 
   // IfOp in ForOp should prevent fusion.
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.if #set(%{{.*}}) {
   // CHECK-NEXT:   }
   // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT: }
   return
 }
 
 // -----
 
 #set = affine_set<(d0) : (d0 - 1 >= 0)>
 
 // CHECK-LABEL: func @should_fuse_if_op_in_loop_nest_not_sandwiched() -> memref<10xf32> {
 func.func @should_fuse_if_op_in_loop_nest_not_sandwiched() -> memref<10xf32> {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %a[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %a[%i1] : memref<10xf32>
     affine.store %v0, %b[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     affine.if #set(%i2) {
       %v0 = affine.load %b[%i2] : memref<10xf32>
     }
   }
 
   // IfOp in ForOp should not prevent fusion if it does not in between the
   // source and dest ForOp ops.
 
   // CHECK:      affine.for
   // CHECK-NEXT:   affine.store
   // CHECK-NEXT:   affine.load
   // CHECK-NEXT:   affine.store
   // CHECK:      affine.for
   // CHECK-NEXT:   affine.if
   // CHECK-NEXT:     affine.load
   // CHECK-NOT:  affine.for
   // CHECK:      return
 
   return %a : memref<10xf32>
 }
 
 // -----
 
 #set = affine_set<(d0) : (d0 - 1 >= 0)>
 
 // CHECK-LABEL: func @should_not_fuse_if_op_in_loop_nest_between_src_and_dest() -> memref<10xf32> {
 func.func @should_not_fuse_if_op_in_loop_nest_between_src_and_dest() -> memref<10xf32> {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %a[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.if #set(%i1) {
       affine.store %cf7, %a[%i1] : memref<10xf32>
     }
   }
   affine.for %i3 = 0 to 10 {
     %v0 = affine.load %a[%i3] : memref<10xf32>
     affine.store %v0, %b[%i3] : memref<10xf32>
   }
   return %b : memref<10xf32>
 
   // IfOp in ForOp which modifies the memref should prevent fusion if it is in
   // between the source and dest ForOp.
 
   // CHECK:      affine.for
   // CHECK-NEXT:   affine.store
   // CHECK:      affine.for
   // CHECK-NEXT:   affine.if
   // CHECK-NEXT:     affine.store
   // CHECK:      affine.for
   // CHECK-NEXT:   affine.load
   // CHECK-NEXT:   affine.store
   // CHECK:      return
 }
 
 // -----
 
 // CHECK-LABEL: func @permute_and_fuse() {
 func.func @permute_and_fuse() {
   %m = memref.alloc() : memref<10x20x30xf32>
 
   %cf7 = arith.constant 7.0 : f32
   affine.for %i0 = 0 to 10 {
     affine.for %i1 = 0 to 20 {
       affine.for %i2 = 0 to 30 {
         affine.store %cf7, %m[%i0, %i1, %i2] : memref<10x20x30xf32>
       }
     }
   }
   affine.for %i3 = 0 to 30 {
     affine.for %i4 = 0 to 10 {
       affine.for %i5 = 0 to 20 {
         %v0 = affine.load %m[%i4, %i5, %i3] : memref<10x20x30xf32>
         "foo"(%v0) : (f32) -> ()
       }
     }
   }
 // CHECK:       affine.for %{{.*}} = 0 to 30 {
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:      affine.for %{{.*}} = 0 to 20 {
 // CHECK-NEXT:        affine.store %{{.*}}, %{{.*}}[0, 0, 0] : memref<1x1x1xf32>
 // CHECK-NEXT:        affine.load %{{.*}}[0, 0, 0] : memref<1x1x1xf32>
 // CHECK-NEXT:        "foo"(%{{.*}}) : (f32) -> ()
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
 
   return
 }
 
 // -----
 
 // CHECK-DAG: [[$MAP0:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> (d0 * 4 + d1)>
 // CHECK-DAG: [[$MAP1:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 floordiv 4)>
 // CHECK-DAG: [[$MAP2:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 4)>
 
 // Reshape from a 64 x f32 to 16 x 4 x f32.
 // CHECK-LABEL: func @fuse_reshape_64_16_4
 func.func @fuse_reshape_64_16_4(%in : memref<64xf32>) {
   %out = memref.alloc() : memref<16x4xf32>
 
   affine.for %i0 = 0 to 64 {
     %v = affine.load %in[%i0] : memref<64xf32>
     affine.store %v, %out[%i0 floordiv 4, %i0 mod 4] : memref<16x4xf32>
   }
 
   affine.for %i1 = 0 to 16 {
     affine.for %i2 = 0 to 4 {
       %w = affine.load %out[%i1, %i2] : memref<16x4xf32>
       "foo"(%w) : (f32) -> ()
     }
   }
   return
   // CHECK:      affine.for %{{.*}} =
   // CHECK-NEXT:   affine.for %{{.*}} =
   // CHECK-NOT:    for
   // CHECK:        }
   // CHECK-NEXT: }
   // CHECK-NEXT: return
 }
 
 // -----
 // CHECK-DAG: [[$MAP0:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 floordiv 4)>
 // CHECK-DAG: [[$MAP1:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 4)>
 // CHECK-DAG: [[$MAP2:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> (d0 * 4 + d1)>
 
 // Reshape a 16x4xf32 to 64xf32.
 // CHECK-LABEL: func @fuse_reshape_16_4_64
 func.func @fuse_reshape_16_4_64() {
   %in = memref.alloc() : memref<16x4xf32>
   %out = memref.alloc() : memref<64xf32>
 
   affine.for %i0 = 0 to 16 {
     affine.for %i1 = 0 to 4 {
       %v = affine.load %in[%i0, %i1] : memref<16x4xf32>
       affine.store %v, %out[4*%i0 + %i1] : memref<64xf32>
     }
   }
 
   affine.for %i2 = 0 to 64 {
     %w = affine.load %out[%i2] : memref<64xf32>
     "foo"(%w) : (f32) -> ()
   }
 // CHECK:       affine.for %{{.*}} = 0 to 64 {
 // CHECK-NEXT:    affine.apply [[$MAP0]](%{{.*}})
 // CHECK-NEXT:    affine.apply [[$MAP1]](%{{.*}})
 // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<16x4xf32>
 // CHECK-NEXT:    affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
 // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
 // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
 // CHECK-NEXT:    "foo"(%{{.*}}) : (f32) -> ()
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
   return
 }
 
 
 // -----
 
 // All three loop nests below (6-d one, 2-d one, 2-d one is fused into a single
 // 2-d loop nest).
 func.func @R6_to_R2_reshape_square() -> memref<64x9xi32> {
   %in = memref.alloc() : memref<2x2x3x3x16x1xi32>
   %out = memref.alloc() : memref<64x9xi32>
   %live_out = memref.alloc() : memref<64x9xi32>
 
   // Initialize input.
   affine.for %i0 = 0 to 2 {
     affine.for %i1 = 0 to 2 {
       affine.for %i2 = 0 to 3 {
         affine.for %i3 = 0 to 3 {
           affine.for %i4 = 0 to 16 {
             affine.for %i5 = 0 to 1 {
               %val = "foo"(%i0, %i1, %i2, %i3, %i4, %i5) : (index, index, index, index, index, index) -> i32
               affine.store %val, %in[%i0, %i1, %i2, %i3, %i4, %i5] : memref<2x2x3x3x16x1xi32>
             }
           }
         }
       }
     }
   }
 
   affine.for %ii = 0 to 64 {
     affine.for %jj = 0 to 9 {
       // Convert output coordinates to linear index.
       %a0 = affine.apply affine_map<(d0, d1) -> (d0 * 9 + d1)> (%ii, %jj)
       %0 = affine.apply affine_map<(d0) -> (d0 floordiv (2 * 3 * 3 * 16 * 1))>(%a0)
       %1 = affine.apply affine_map<(d0) -> ((d0 mod 288) floordiv (3 * 3 * 16 * 1))>(%a0)
       %2 = affine.apply affine_map<(d0) -> (((d0 mod 288) mod 144) floordiv (3 * 16 * 1))>(%a0)
       %3 = affine.apply affine_map<(d0) -> ((((d0 mod 288) mod 144) mod 48) floordiv (16 * 1))>(%a0)
       %4 = affine.apply affine_map<(d0) -> ((((d0 mod 288) mod 144) mod 48) mod 16)>(%a0)
       %5 = affine.apply affine_map<(d0) -> (((((d0 mod 144) mod 144) mod 48) mod 16) mod 1)>(%a0)
       %v = affine.load %in[%0, %1, %2, %3, %4, %5] : memref<2x2x3x3x16x1xi32>
       affine.store %v, %out[%ii, %jj] : memref<64x9xi32>
     }
   }
 
   affine.for %i = 0 to 64 {
     affine.for %j = 0 to 9 {
       %a = affine.load %out[%i, %j] : memref<64x9xi32>
       %b = arith.muli %a, %a : i32
       affine.store %b, %live_out[%i, %j] : memref<64x9xi32>
     }
   }
   return %live_out : memref<64x9xi32>
 }
 // Everything above is fused to a single 2-d loop nest, and the 6-d tensor %in
 // is eliminated if -memref-dataflow-opt is also supplied.
 //
 // CHECK-DAG: [[$MAP0:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> ((d0 * 9 + d1) floordiv 288)>
 // CHECK-DAG: [[$MAP1:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> (((d0 * 9 + d1) mod 288) floordiv 144)>
 // CHECK-DAG: [[$MAP2:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> (((d0 * 9 + d1) mod 144) floordiv 48)>
 // CHECK-DAG: [[$MAP3:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> (((d0 * 9 + d1) mod 48) floordiv 16)>
 // CHECK-DAG: [[$MAP4:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> ((d0 * 9 + d1) mod 16)>
 // CHECK-DAG: [[$MAP11:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> (d0 * 9 + d1)>
 // CHECK-DAG: [[$MAP12:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 floordiv 288)>
 // CHECK-DAG: [[$MAP13:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> ((d0 mod 288) floordiv 144)>
 // CHECK-DAG: [[$MAP14:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> ((d0 mod 144) floordiv 48)>
 // CHECK-DAG: [[$MAP15:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> ((d0 mod 48) floordiv 16)>
 // CHECK-DAG: [[$MAP16:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 16)>
 // CHECK-DAG: [[$MAP17:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (0)>
 
 //
 // CHECK-LABEL: func @R6_to_R2_reshape
 // CHECK:       memref.alloc() : memref<1x2x3x3x16x1xi32>
 // CHECK:       memref.alloc() : memref<1x1xi32>
 // CHECK:       memref.alloc() : memref<64x9xi32>
 // CHECK-NEXT:  affine.for %{{.*}} = 0 to 64 {
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 9 {
 // CHECK-NEXT:      affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP1]](%{{.*}}, %{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP3]](%{{.*}}, %{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP4]](%{{.*}}, %{{.*}})
 // CHECK-NEXT:      "foo"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index, index, index, index) -> i32
 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, ((%{{.*}} * 9 + %{{.*}}) mod 288) floordiv 144, ((%{{.*}} * 9 + %{{.*}}) mod 144) floordiv 48, ((%{{.*}} * 9 + %{{.*}}) mod 48) floordiv 16, (%{{.*}} * 9 + %{{.*}}) mod 16, 0] : memref<1x2x3x3x16x1xi32>
 // CHECK-NEXT:      affine.apply [[$MAP11]](%{{.*}}, %{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP12]](%{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP13]](%{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP14]](%{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP15]](%{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP16]](%{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP17]](%{{.*}})
 // CHECK-NEXT:      affine.load %{{.*}}[0, ((%{{.*}} * 9 + %{{.*}}) mod 288) floordiv 144, ((%{{.*}} * 9 + %{{.*}}) mod 144) floordiv 48, ((%{{.*}} * 9 + %{{.*}}) mod 48) floordiv 16, (%{{.*}} * 9 + %{{.*}}) mod 16, 0] : memref<1x2x3x3x16x1xi32>
 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xi32>
 // CHECK-NEXT:      affine.load %{{.*}}[0, 0] : memref<1x1xi32>
 // CHECK-NEXT:      arith.muli %{{.*}}, %{{.*}} : i32
 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x9xi32>
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return %{{.*}} : memref<64x9xi32>
 
 // -----
 
 // CHECK-LABEL: func @fuse_symbolic_bounds
 func.func @fuse_symbolic_bounds(%M : index, %N : index) {
   %N_plus_5 = affine.apply affine_map<(d0) -> (d0 + 5)>(%N)
   %m = memref.alloc(%M, %N_plus_5) : memref<? x ? x f32>
 
   %c0 = arith.constant 0.0 : f32
   %s = arith.constant 5 : index
 
   affine.for %i0 = 0 to %M {
     affine.for %i1 = 0 to affine_map<(d0) -> (d0 + 5)> (%N) {
       affine.store %c0, %m[%i0, %i1] : memref<? x ? x f32>
     }
   }
 
   affine.for %i2 = 0 to %M {
     affine.for %i3 = 0 to %N {
       %v = affine.load %m[%i2, %i3 + symbol(%s)] : memref<? x ? x f32>
     }
   }
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_reduction_at_depth_of_one
 func.func @should_fuse_reduction_at_depth_of_one() {
   %a = memref.alloc() : memref<10x100xf32>
   %b = memref.alloc() : memref<10xf32>
 
   affine.for %i0 = 0 to 10 {
     affine.for %i1 = 0 to 100 {
       %v0 = affine.load %b[%i0] : memref<10xf32>
       %v1 = affine.load %a[%i0, %i1] : memref<10x100xf32>
       %v2 = "maxf"(%v0, %v1) : (f32, f32) -> f32
       affine.store %v2, %b[%i0] : memref<10xf32>
     }
   }
   affine.for %i2 = 0 to 10 {
     affine.for %i3 = 0 to 100 {
       %v3 = affine.load %b[%i2] : memref<10xf32>
       %v4 = affine.load %a[%i2, %i3] : memref<10x100xf32>
       %v5 = arith.subf %v4, %v3 : f32
       affine.store %v5, %b[%i2] : memref<10xf32>
     }
   }
   // This test should fuse the src reduction loop at depth 1 in the destination
   // loop nest, which improves locality and enables subsequence passes to
   // decrease the reduction memref size and possibly place it in a faster
   // memory space.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 100 {
   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x100xf32>
   // CHECK-NEXT:      "maxf"(%{{.*}}, %{{.*}}) : (f32, f32) -> f32
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 100 {
   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x100xf32>
   // CHECK-NEXT:      arith.subf %{{.*}}, %{{.*}} : f32
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_at_src_depth1_and_dst_depth1
 func.func @should_fuse_at_src_depth1_and_dst_depth1() {
   %a = memref.alloc() : memref<100x16xf32>
   %b = memref.alloc() : memref<100x16xf32>
 
   affine.for %i0 = 0 to 100 {
     affine.for %i1 = 0 to 16 {
       %v0 = affine.load %a[%i0, %i1] : memref<100x16xf32>
       "op0"(%v0) : (f32) -> ()
     }
     affine.for %i2 = 0 to 16 {
       %v1 = "op1"() : () -> (f32)
       affine.store %v1, %b[%i0, %i2] : memref<100x16xf32>
     }
   }
 
   affine.for %i3 = 0 to 100 {
     affine.for %i4 = 0 to 16 {
       %v2 = affine.load %b[%i3, %i4] : memref<100x16xf32>
       "op2"(%v2) : (f32) -> ()
     }
   }
   // We can slice iterations of the '%i0' and '%i1' loops in the source
   // loop nest, but slicing at depth 2 and inserting the slice in the
   // destination loop nest at depth2 causes extra computation. Instead,
   // the fusion algorithm should detect that the source loop should be sliced
   // at depth 1 and the slice should be inserted at depth 1.
   // CHECK:       affine.for %{{.*}} = 0 to 100 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<100x16xf32>
   // CHECK-NEXT:      "op0"(%{{.*}}) : (f32) -> ()
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
   // CHECK-NEXT:      %{{.*}} = "op1"() : () -> f32
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
   // CHECK-NEXT:      affine.load %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
   // CHECK-NEXT:      "op2"(%{{.*}}) : (f32) -> ()
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 // CHECK: [[$MAP0:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 * 10 + d1)>
 
 // CHECK-LABEL: func @should_fuse_src_depth1_at_dst_depth2
 func.func @should_fuse_src_depth1_at_dst_depth2() {
   %a = memref.alloc() : memref<100xf32>
   %c0 = arith.constant 0.0 : f32
 
   affine.for %i0 = 0 to 100 {
     affine.store %c0, %a[%i0] : memref<100xf32>
   }
 
   affine.for %i1 = 0 to 10 {
     affine.for %i2 = 0 to 10 {
       %a0 = affine.apply affine_map<(d0, d1) -> (d0 * 10 + d1)> (%i1, %i2)
       %v0 = affine.load %a[%a0] : memref<100xf32>
     }
   }
   // The source loop nest slice loop bound is a function of both destination
   // loop IVs, so we should slice at depth 1 and insert the slice at depth 2.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:      affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:      affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
   // CHECK-NEXT:      affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @fusion_at_depth0_not_currently_supported
 func.func @fusion_at_depth0_not_currently_supported() {
   %0 = memref.alloc() : memref<10xf32>
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %0[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %1 = affine.load %0[%c0] : memref<10xf32>
   }
   // NOTE: Should shrink memref size to 1 element access by load in dst loop
   // nest, and make the store in the slice store to the same element.
   // CHECK-DAG:   memref.alloc() : memref<1xf32>
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_deep_loop_nests
 func.func @should_fuse_deep_loop_nests() {
   %0 = memref.alloc() : memref<2x2x3x3x16x10xf32, 2>
   %1 = memref.alloc() : memref<2x2x3x3x16x10xf32, 2>
   %2 = memref.alloc() : memref<3x3x3x3x16x10xf32, 2>
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c1_0 = arith.constant 1 : index
   %cst = arith.constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 2 {
     affine.for %i1 = 0 to 2 {
       affine.for %i2 = 0 to 3 {
         affine.for %i3 = 0 to 3 {
           affine.for %i4 = 0 to 16 {
             affine.for %i5 = 0 to 10 {
               %3 = affine.load %0[%i0, %i1, %i2, %i3, %i4, %i5]
                 : memref<2x2x3x3x16x10xf32, 2>
             }
           }
           affine.for %i6 = 0 to 16 {
             affine.for %i7 = 0 to 10 {
               affine.store %cst, %1[%i0, %i1, %i2, %i3, %i6, %i7]
                 : memref<2x2x3x3x16x10xf32, 2>
             }
           }
         }
       }
     }
   }
   affine.for %i8 = 0 to 3 {
     affine.for %i9 = 0 to 3 {
       affine.for %i10 = 0 to 2 {
         affine.for %i11 = 0 to 2 {
           affine.for %i12 = 0 to 3 {
             affine.for %i13 = 0 to 3 {
               affine.for %i14 = 0 to 2 {
                 affine.for %i15 = 0 to 2 {
                   affine.for %i16 = 0 to 16 {
                     affine.for %i17 = 0 to 10 {
                       %5 = affine.load %0[%i14, %i15, %i12, %i13, %i16, %i17]
                         : memref<2x2x3x3x16x10xf32, 2>
                     }
                   }
                   affine.for %i18 = 0 to 16 {
                     affine.for %i19 = 0 to 10 {
                       %6 = affine.load %1[%i10, %i11, %i8, %i9, %i18, %i19]
                         : memref<2x2x3x3x16x10xf32, 2>
                     }
                   }
                 }
               }
             }
           }
         }
       }
     }
   }
 // The first four loops of the source loop nest can be sliced with iteration
 // bounds which are a function of the first four loops of destination loop nest,
 // where the destination loops nests have been interchanged.
 
 // CHECK-DAG:   memref.alloc() : memref<1x1x1x1x16x10xf32, 2>
 // CHECK:       affine.for %{{.*}} = 0 to 3 {
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 3 {
 // CHECK-NEXT:      affine.for %{{.*}} = 0 to 2 {
 // CHECK-NEXT:        affine.for %{{.*}} = 0 to 2 {
 // CHECK-NEXT:          affine.for %{{.*}} = 0 to 3 {
 // CHECK-NEXT:            affine.for %{{.*}} = 0 to 3 {
 // CHECK-NEXT:              affine.for %{{.*}} = 0 to 16 {
 // CHECK-NEXT:                affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:                  affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<2x2x3x3x16x10xf32, 2>
 // CHECK-NEXT:                }
 // CHECK-NEXT:              }
 // CHECK-NEXT:              affine.for %{{.*}} = 0 to 16 {
 // CHECK-NEXT:                affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:                  affine.store %{{.*}}, %{{.*}}[0, 0, 0, 0, %{{.*}}, %{{.*}}] : memref<1x1x1x1x16x10xf32, 2>
 // CHECK-NEXT:                }
 // CHECK-NEXT:              }
 // CHECK-NEXT:              affine.for %{{.*}} = 0 to 2 {
 // CHECK-NEXT:                affine.for %{{.*}} = 0 to 2 {
 // CHECK-NEXT:                  affine.for %{{.*}} = 0 to 16 {
 // CHECK-NEXT:                    affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:                      affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<2x2x3x3x16x10xf32, 2>
 // CHECK-NEXT:                    }
 // CHECK-NEXT:                  }
 // CHECK-NEXT:                  affine.for %{{.*}} = 0 to 16 {
 // CHECK-NEXT:                    affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:                      affine.load %{{.*}}[0, 0, 0, 0, %{{.*}}, %{{.*}}] : memref<1x1x1x1x16x10xf32, 2>
 // CHECK-NEXT:                    }
 // CHECK-NEXT:                  }
 // CHECK-NEXT:                }
 // CHECK-NEXT:              }
 // CHECK-NEXT:            }
 // CHECK-NEXT:          }
 // CHECK-NEXT:        }
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_at_depth1_and_reduce_slice_trip_count
 func.func @should_fuse_at_depth1_and_reduce_slice_trip_count() {
   %a = memref.alloc() : memref<4x256xf32>
   %b = memref.alloc() : memref<4x256xf32>
 
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
 
   affine.for %i0 = 0 to 4 {
     affine.for %i1 = 0 to 256 {
       %v0 = affine.load %b[%i0, %i1] : memref<4x256xf32>
     }
     affine.for %i2 = 0 to 256 {
       affine.store %cf0, %a[%i0, %i2] : memref<4x256xf32>
     }
   }
 
   affine.for %d0 = 0 to 4 {
     affine.for %d1 = 0 to 16 {
       %v1 = affine.load %a[%d0, %d1] : memref<4x256xf32>
     }
   }
   // The cost of fusing at depth 2 is greater than the cost of fusing at depth 1
   // for two reasons:
   // 1) Inserting the unsliceable src loop %i1 to a higher depth removes
   //    redundant computation and reduces costs.
   // 2) Inserting the sliceable src loop %i2 at depth 1, we can still reduce
   //    its trip count to 16 (from 256) reducing costs.
   // NOTE: the size of the private memref created for the fused loop nest
   // is reduced from the original shape from 4x256 to 4x16 because of the
   // data accessed by the load.
   // CHECK-DAG:   memref.alloc() : memref<1x16xf32>
   // CHECK:       affine.for %{{.*}} = 0 to 4 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 256 {
   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<4x256xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 16 {
   // CHECK-NEXT:      affine.load %{{.*}}[0, %{{.*}}] : memref<1x16xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_at_depth1_with_trip_count_20
 func.func @should_fuse_at_depth1_with_trip_count_20() {
   %a = memref.alloc() : memref<100xf32>
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
 
   affine.for %i0 = 0 to 100 {
     affine.store %cf0, %a[%i0]: memref<100xf32>
   }
 
   affine.for %i1 = 0 to 5 {
     affine.for %i2 = 0 to 10 {
       %v0 = affine.load %a[%i2]: memref<100xf32>
     }
     affine.for %i3 = 0 to 10 {
       affine.for %i4 = 0 to 20 {
         %v1 = affine.load %a[%i4]: memref<100xf32>
       }
     }
   }
   // NOTE: The size of the private memref created for fusion is shrunk to 20xf32
   // CHECK-DAG:   memref.alloc() : memref<20xf32>
   // CHECK:       affine.for %{{.*}} = 0 to 5 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 20 {
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<20xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}] : memref<20xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 20 {
   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}] : memref<20xf32>
   // CHECK-NEXT:      }
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_at_depth1_with_trip_count_19
 func.func @should_fuse_at_depth1_with_trip_count_19() {
   %a = memref.alloc() : memref<100xf32>
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
 
   affine.for %i0 = 0 to 100 {
     affine.store %cf0, %a[%i0]: memref<100xf32>
   }
 
   affine.for %i1 = 0 to 5 {
     affine.for %i2 = 0 to 19 {
       %v0 = affine.load %a[%i2]: memref<100xf32>
     }
     affine.for %i3 = 0 to 10 {
       affine.for %i4 = 0 to 10 {
         %v1 = affine.load %a[%i4]: memref<100xf32>
       }
     }
   }
   // NOTE: The size of the private memref created for fusion is shrunk to 19xf32
   // CHECK-DAG:   memref.alloc() : memref<19xf32>
   // CHECK:       affine.for %{{.*}} = 0 to 5 {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 19 {
   // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<19xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 19 {
   // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}] : memref<19xf32>
   // CHECK-NEXT:    }
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:        affine.load %{{.*}}[%{{.*}}] : memref<19xf32>
   // CHECK-NEXT:      }
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_with_private_memref() {
 func.func @should_fuse_with_private_memref() {
   %m = memref.alloc() : memref<100xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 100 {
     affine.store %cf7, %m[%i0] : memref<100xf32>
   }
   affine.for %i1 = 0 to 17 {
     %v0 = affine.load %m[%i1] : memref<100xf32>
   }
   affine.for %i2 = 0 to 82 {
     %v1 = affine.load %m[%i2] : memref<100xf32>
   }
   // Should create a new private memref.
   // CHECK-DAG:  memref.alloc() : memref<1xf32>
   // CHECK:      affine.for %{{.*}} = 0 to 17 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_live_out_arg_but_preserve_src_loop(%{{.*}}: memref<10xf32>) {
 func.func @should_fuse_live_out_arg_but_preserve_src_loop(%arg0: memref<10xf32>) {
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %arg0[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 9 {
     %v0 = affine.load %arg0[%i1] : memref<10xf32>
   }
   // This tests that the loop nest '%i0' should not be removed after fusion
   // because it writes to memref argument '%arg0', and its read region
   // does not cover its write region (so fusion would shrink the write region
   // in the fused loop nest, so complete live out data region would not
   // be written).
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}} : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 9 {
   // CHECK-NEXT:    affine.store %{{.*}} : memref<1xf32>
   // CHECK-NEXT:    affine.load %{{.*}} : memref<1xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_live_out_arg(%{{.*}}: memref<10xf32>) {
 func.func @should_fuse_live_out_arg(%arg0: memref<10xf32>) {
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %arg0[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %arg0[%i1] : memref<10xf32>
   }
   // The read/write regions for memref '%{{.*}}' are the same for both
   // loops, so they should fuse.
 
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_escaping_memref_but_preserve_src_loop() -> memref<10xf32>
 func.func @should_fuse_escaping_memref_but_preserve_src_loop() -> memref<10xf32> {
   %cf7 = arith.constant 7.0 : f32
   %m = memref.alloc() : memref<10xf32>
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 9 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // This tests that the loop nest '%i0' should not be removed after fusion
   // because it writes to memref '%m', which is returned by the function, and
   // the '%i1' memory region does not cover '%i0' memory region.
 
   // CHECK-DAG:   memref.alloc() : memref<1xf32>
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}} : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 9 {
   // CHECK-NEXT:    affine.store %{{.*}} : memref<1xf32>
   // CHECK-NEXT:    affine.load %{{.*}} : memref<1xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return %{{.*}} : memref<10xf32>
   return %m : memref<10xf32>
 }
 // -----
 
 // This should fuse with the %in becoming a 1x1x1.
 func.func @R3_to_R2_reshape() {
   %in = memref.alloc() : memref<2x3x16xi32>
 
   %c0 = arith.constant 0 : index
 
   affine.for %i0 = 0 to 2 {
     affine.for %i1 = 0 to 3 {
       affine.for %i2 = 0 to 16 {
         %val = "foo"(%i0, %i1, %i2) : (index, index, index) -> i32
         affine.store %val, %in[%i0, %i1, %i2] : memref<2x3x16xi32>
       }
     }
   }
 
   affine.for %ii = 0 to 32 {
     affine.for %jj = 0 to 3 {
       %a0 = affine.apply affine_map<(d0, d1) -> (d0 * 3 + d1)> (%ii, %jj)
       %idx = affine.apply affine_map<(d0) -> (d0 floordiv (3 * 16))> (%a0)
       %v = affine.load %in[%idx, %jj, %c0]
         : memref<2x3x16xi32>
     }
   }
   return
 }
 // CHECK-DAG: [[$MAP0:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> ((d0 * 3 + d1) floordiv 48)>
 // CHECK-DAG: [[$MAP1:#map[0-9a-zA-Z_]*]] = affine_map<(d0, d1) -> (d0 * 3 + d1)>
 // CHECK-DAG: [[$MAP2:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 floordiv 48)>
 
 // CHECK-LABEL: func @R3_to_R2_reshape()
 // CHECK-DAG:    memref.alloc() : memref<1x1x1xi32>
 // CHECK:        affine.for %{{.*}} = 0 to 32 {
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 3 {
 // CHECK-NEXT:      affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
 // CHECK-NEXT:      "foo"(%{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index) -> i32
 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[0, 0, 0] : memref<1x1x1xi32>
 // CHECK-NEXT:      affine.apply [[$MAP1]](%{{.*}}, %{{.*}})
 // CHECK-NEXT:      affine.apply [[$MAP2]](%{{.*}})
 // CHECK-NEXT:      affine.load %{{.*}}[0, 0, 0] : memref<1x1x1xi32>
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
 
 // -----
 
 func.func @should_fuse_multi_output_producer() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %a[%i0] : memref<10xf32>
     affine.store %cf7, %b[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     %v0 = affine.load %a[%i1] : memref<10xf32>
     %v1 = affine.load %b[%i1] : memref<10xf32>
   }
 
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @fusion_preventing_deps_on_middle_loop() {
 func.func @fusion_preventing_deps_on_middle_loop() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %c = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     %v0 = affine.load %a[%i0] : memref<10xf32>
     affine.store %v0, %b[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.store %cf7, %a[%i1] : memref<10xf32>
     %v1 = affine.load %c[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     %v2 = affine.load %b[%i2] : memref<10xf32>
     affine.store %v2, %c[%i2] : memref<10xf32>
   }
   // Loops '%i0' and '%i2' cannot fuse along producer/consumer edge on memref
   // '%b', because of the WAR dep from '%i0' to '%i1' on memref '%a' and
   // because of the WAR dep from '%i1' to '%i2' on memref '%c'.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_and_move_to_preserve_war_dep() {
 func.func @should_fuse_and_move_to_preserve_war_dep() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %c = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     %v0 = affine.load %b[%i0] : memref<10xf32>
     affine.store %v0, %a[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 3 {
     %v2 = affine.load %c[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 5 {
     affine.store %cf7, %b[%i2] : memref<10xf32>
   }
   affine.for %i3 = 0 to 10 {
     %v1 = affine.load %a[%i3] : memref<10xf32>
     affine.store %cf7, %c[%i3] : memref<10xf32>
   }
 
   // Dependence graph:
   //
   //         %i0 ---------
   //               |     |
   //     --- %i1   | %b  | %a
   //     |         |     |
   //  %c |   %i2 <--     |
   //     |               |
   //     --> %i3 <--------
   //
   // It is possible to fuse loop '%i0' into '%i3' and preserve dependences
   // if the fused loop nest is inserted between loops '%i1' and '%i2'.
 
   // CHECK-DAG:   memref.alloc() : memref<1xf32>
   // CHECK:       affine.for %{{.*}} = 0 to 3 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 5 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @fusion_preventing_dep_on_constant() {
 func.func @fusion_preventing_dep_on_constant() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %c = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     %v0 = affine.load %b[%i0] : memref<10xf32>
     affine.store %cf7, %a[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.store %cf7, %b[%i1] : memref<10xf32>
   }
   %cf11 = arith.constant 11.0 : f32
   affine.for %i2 = 0 to 10 {
     %v2 = affine.load %a[%i2] : memref<10xf32>
     affine.store %cf11, %c[%i2] : memref<10xf32>
   }
   // Loops '%i0' and '%i2' cannot fuse along producer/consumer edge on memref
   // '%a', because of the WAR dep from '%i0' to '%i1' on memref '%b' and
   // because of the SSA value dep from '%cf11' def to use in '%i2'.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  %{{.*}} = arith.constant 1.100000e+01 : f32
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_and_preserve_dep_on_constant() {
 func.func @should_fuse_and_preserve_dep_on_constant() {
   %a = memref.alloc() : memref<10xf32>
   %b = memref.alloc() : memref<10xf32>
   %c = memref.alloc() : memref<10xf32>
 
   %cf7 = arith.constant 7.0 : f32
   %cf11 = arith.constant 11.0 : f32
   affine.for %i0 = 0 to 10 {
     %v0 = affine.load %b[%i0] : memref<10xf32>
     affine.store %cf7, %a[%i0] : memref<10xf32>
   }
   affine.for %i1 = 0 to 10 {
     affine.store %cf7, %b[%i1] : memref<10xf32>
   }
   affine.for %i2 = 0 to 10 {
     %v2 = affine.load %a[%i2] : memref<10xf32>
     affine.store %cf11, %c[%i2] : memref<10xf32>
   }
 
   // Loops '%i0' and '%i2' can fuse along producer/consumer edge on memref
   // '%a', and preserve the WAR dep from '%i0' to '%i1' on memref '%b', and
   // the SSA value dep from '%cf11' def to use in '%i2'.
 
   // CHECK:       arith.constant 1.100000e+01 : f32
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // Add further tests in mlir/test/Transforms/loop-fusion-4.mlir
 

mlir/test/Dialect/Affine/pipeline-data-transfer.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -affine-pipeline-data-transfer | FileCheck %s
 
 // -----
 
 // CHECK-DAG: [[$MOD_2:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 2)>
 // CHECK-DAG: [[$MAP_MINUS_1:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 - 1)>
 
 // CHECK-LABEL: func @loop_nest_dma() {
 func.func @loop_nest_dma() {
 
   %A = memref.alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 0>
   %Ah = memref.alloc() : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
 
   %tag = memref.alloc() : memref<1 x f32>
 
   %zero = arith.constant 0 : index
   %num_elts = arith.constant 32 : index
 
   affine.for %i = 0 to 8 {
     affine.dma_start %A[%i], %Ah[%i], %tag[%zero], %num_elts : memref<256 x f32>, memref<32 x f32, 1>, memref<1 x f32>
     affine.dma_wait %tag[%zero], %num_elts : memref<1 x f32>
     %v = affine.load %Ah[%i] : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
     %r = "compute"(%v) : (f32) -> (f32)
     affine.store %r, %Ah[%i] : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
     affine.for %j = 0 to 32 {
       "do_more_compute"(%i, %j) : (index, index) -> ()
     }
   }
   memref.dealloc %tag : memref<1 x f32>
   memref.dealloc %Ah : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
   return
 }
 // CHECK:       %{{.*}} = memref.alloc() : memref<256xf32>
 // CHECK:       %{{.*}} = memref.alloc() : memref<2x32xf32, 1>
 // CHECK-NEXT:  %{{.*}} = memref.alloc() : memref<2x1xf32>
 // CHECK-NEXT:  affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}} mod 2, %{{.*}}], %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<256xf32>, memref<2x32xf32, 1>, memref<2x1xf32>
 // CHECK-NEXT:  affine.for %{{.*}} = 1 to 8 {
 // CHECK-NEXT:    affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}} mod 2, %{{.*}}], %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<256xf32>, memref<2x32xf32, 1>, memref<2x1xf32>
 // CHECK-NEXT:    affine.apply [[$MAP_MINUS_1]](%{{.*}})
 // CHECK-NEXT:    affine.apply [[$MOD_2]](%{{.*}})
 // CHECK-NEXT:    affine.apply [[$MOD_2]](%{{.*}})
 // CHECK-NEXT:    affine.dma_wait %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<2x1xf32>
 // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}} mod 2, %{{.*}}] : memref<2x32xf32, 1>
 // CHECK-NEXT:    "compute"(%{{.*}}) : (f32) -> f32
 // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}} mod 2, %{{.*}}] : memref<2x32xf32, 1>
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 32 {
 // CHECK-NEXT:      "do_more_compute"(%{{.*}}, %{{.*}}) : (index, index) -> ()
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  affine.apply [[$MAP_MINUS_1]](%{{.*}})
 // CHECK-NEXT:  affine.apply [[$MOD_2]](%{{.*}})
 // CHECK-NEXT:  affine.apply [[$MOD_2]](%{{.*}})
 // CHECK-NEXT:  affine.dma_wait %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<2x1xf32>
 // CHECK-NEXT:  affine.load %{{.*}}[%{{.*}} mod 2, %{{.*}}] : memref<2x32xf32, 1>
 // CHECK-NEXT:  "compute"(%{{.*}}) : (f32) -> f32
 // CHECK-NEXT:  affine.store %{{.*}}, %{{.*}}[%{{.*}} mod 2, %{{.*}}] : memref<2x32xf32, 1>
 // CHECK-NEXT:  affine.for %{{.*}} = 0 to 32 {
 // CHECK-NEXT:    "do_more_compute"(%{{.*}}, %{{.*}}) : (index, index) -> ()
 // CHECK-NEXT:  }
 // CHECK-NEXT:  memref.dealloc %{{.*}} : memref<2x1xf32>
 // CHECK-NEXT:  memref.dealloc %{{.*}} : memref<2x32xf32, 1>
 // CHECK-NEXT:  return
 // CHECK-NEXT:}
 
 // -----
 
 // CHECK-DAG: [[$FLOOR_MOD_2:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> ((d0 floordiv 4) mod 2)>
 // CHECK-DAG: [[$REMAP_SHIFT_MINUS_4:#map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 - 4)>
 
 // CHECK-LABEL: @loop_step
 func.func @loop_step(%arg0: memref<512xf32>,
                   %arg1: memref<512xf32>) {
   %c0 = arith.constant 0 : index
   %c4 = arith.constant 4 : index
   affine.for %i0 = 0 to 512 step 4 {
     %1 = memref.alloc() : memref<4xf32, 1>
     %2 = memref.alloc() : memref<1xi32>
     affine.dma_start %arg0[%i0], %1[%c0], %2[%c0], %c4,
               : memref<512xf32>, memref<4xf32, 1>, memref<1xi32>
     affine.dma_wait %2[%c0], %c4 : memref<1xi32>
     "compute"(%i0) : (index) -> ()
     memref.dealloc %2 : memref<1xi32>
     memref.dealloc %1 : memref<4xf32, 1>
   }
   return
 }
 // CHECK:        [[BUF:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x4xf32, 1>
 // CHECK:        [[TAG:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x1xi32>
 // CHECK-NEXT:   affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[(%{{.*}} floordiv 4) mod 2, 0], [[TAG]][(%{{.*}} floordiv 4) mod 2, 0], %{{.*}} : memref<512xf32>, memref<2x4xf32, 1>, memref<2x1xi32>
 // CHECK-NEXT:   affine.for %{{.*}} = 4 to 512 step 4 {
 // CHECK-NEXT:     affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[(%{{.*}} floordiv 4) mod 2, 0], [[TAG]][(%{{.*}} floordiv 4) mod 2, 0], %{{.*}} : memref<512xf32>, memref<2x4xf32, 1>, memref<2x1xi32>
 // CHECK-NEXT:     affine.apply [[$REMAP_SHIFT_MINUS_4]](%{{.*}})
 // CHECK-NEXT:     affine.apply [[$FLOOR_MOD_2]](%{{.*}})
 // CHECK:          affine.dma_wait [[TAG]][(%{{.*}} floordiv 4) mod 2, 0], %{{.*}} : memref<2x1xi32>
 // CHECK-NEXT:     "compute"(%{{.*}}) : (index) -> ()
 // CHECK-NEXT:   }
 // CHECK-NEXT:   [[SHIFTED:%[0-9a-zA-Z_]+]] = affine.apply [[$REMAP_SHIFT_MINUS_4]](%{{.*}})
 // CHECK-NEXT:   %{{.*}} = affine.apply [[$FLOOR_MOD_2]]([[SHIFTED]])
 // CHECK:        affine.dma_wait [[TAG]][(%{{.*}} floordiv 4) mod 2, 0], %{{.*}} : memref<2x1xi32>
 // CHECK-NEXT:   "compute"(%{{.*}}) : (index) -> ()
 // CHECK-NEXT:   memref.dealloc [[TAG]] : memref<2x1xi32>
 // CHECK-NEXT:   memref.dealloc [[BUF]] : memref<2x4xf32, 1>
 // CHECK-NEXT:   return
 // CHECK-NEXT: }
 
 // -----
 
 #map1 = affine_map<(d0, d1) -> ((d0 * 2048 + d1 * 256) floordiv 32)>
 #map2 = affine_map<(d0) -> ((d0 * 2048) floordiv 32)>
 // CHECK-LABEL: func @loop_dma_nested(%{{.*}}: memref<512x32xvector<8xf32>
 func.func @loop_dma_nested(%arg0: memref<512x32xvector<8xf32>>, %arg1: memref<512x32xvector<8xf32>>, %arg2: memref<512x32xvector<8xf32>>) {
   %num_elts = arith.constant 256 : index
   %c0 = arith.constant 0 : index
   %0 = memref.alloc() : memref<64x4xvector<8xf32>, 2>
   %1 = memref.alloc() : memref<64x4xvector<8xf32>, 2>
   %2 = memref.alloc() : memref<64x4xvector<8xf32>, 2>
   %3 = memref.alloc() : memref<2xi32>
   %4 = memref.alloc() : memref<2xi32>
   %5 = memref.alloc() : memref<2xi32>
   // Prologue for DMA overlap on arg2.
   // CHECK-DAG: [[BUF_ARG2:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x64x4xvector<8xf32>, 2>
   // CHECK-DAG: [[TAG_ARG2:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x2xi32>
   // CHECK: affine.dma_start %{{.*}}[
   // CHECK: affine.for %{{.*}} = 1 to 8 {
   affine.for %i0 = 0 to 8 {
     %6 = affine.apply #map2(%i0)
     affine.dma_start %arg2[%6, %c0], %2[%c0, %c0], %5[%c0], %num_elts : memref<512x32xvector<8xf32>>, memref<64x4xvector<8xf32>, 2>, memref<2xi32>
     affine.dma_wait %5[%c0], %num_elts : memref<2xi32>
     // Steady state for DMA overlap on arg2
     // CHECK: affine.dma_start %{{.*}}[
     // CHECK: affine.dma_wait [[TAG_ARG2]]
     // Prologue for DMA overlap on arg0, arg1 nested within i0
     // CHECK: [[BUF_ARG0:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x64x4xvector<8xf32>, 2>
     // CHECK: [[BUF_ARG1:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x64x4xvector<8xf32>, 2>
     // CHECK: [[TAG_ARG0:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x2xi32>
     // CHECK: [[TAG_ARG1:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x2xi32>
     // CHECK: affine.dma_start %{{.*}}[
     // CHECK: affine.dma_start %{{.*}}[
     // CHECK-NEXT: affine.for %{{.*}} = 1 to 8 {
     affine.for %i1 = 0 to 8 {
       %7 = affine.apply #map1(%i0, %i1)
       %8 = affine.apply #map2(%i1)
       affine.dma_start %arg0[%7, %c0], %0[%c0, %c0], %3[%c0], %num_elts : memref<512x32xvector<8xf32>>, memref<64x4xvector<8xf32>, 2>, memref<2xi32>
       affine.dma_start %arg1[%8, %c0], %1[%c0, %c0], %4[%c0], %num_elts : memref<512x32xvector<8xf32>>, memref<64x4xvector<8xf32>, 2>, memref<2xi32>
       affine.dma_wait %3[%c0], %num_elts : memref<2xi32>
       affine.dma_wait %4[%c0], %num_elts : memref<2xi32>
       // Steady state for DMA overlap on arg0, arg1
       // CHECK: affine.dma_start %{{.*}}[
       // CHECK: affine.dma_start %{{.*}}[
       // CHECK: affine.dma_wait [[TAG_ARG0]]
       // CHECK: affine.dma_wait [[TAG_ARG1]]
       // CHECK-NEXT: affine.for %{{.*}} = 0 to 4 {
       affine.for %i2 = 0 to 4 {
         "foo"() : () -> ()
       }
     }
     // epilogue for arg0, arg1
     // CHECK: affine.dma_wait [[TAG_ARG0]]
     // CHECK: affine.dma_wait [[TAG_ARG1]]
     // CHECK-DAG:    memref.dealloc [[TAG_ARG1]] : memref<2x2xi32>
     // CHECK-DAG:    memref.dealloc [[TAG_ARG0]] : memref<2x2xi32>
     // CHECK-DAG:    memref.dealloc [[BUF_ARG1]] : memref<2x64x4xvector<8xf32>, 2>
     // CHECK-DAG:    memref.dealloc [[BUF_ARG0]] : memref<2x64x4xvector<8xf32>, 2>
   // epilogue for DMA overlap on %arg2
   // CHECK:  affine.dma_wait [[TAG_ARG2]]
   // Within the epilogue for arg2's DMA, we have the DMAs on %arg1, %arg2 nested.
   // CHECK: [[BUF_ARG0_NESTED:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x64x4xvector<8xf32>, 2>
   // CHECK: [[BUF_ARG1_NESTED:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x64x4xvector<8xf32>, 2>
   // CHECK: [[TAG_ARG0_NESTED:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x2xi32>
   // CHECK: [[TAG_ARG1_NESTED:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x2xi32>
   // CHECK:  affine.dma_start %{{.*}}[
   // CHECK:  affine.dma_start %{{.*}}[
   // CHECK:  affine.for %{{.*}} = 1 to 8 {
   // CHECK:    affine.dma_start %{{.*}}[
   // CHECK:    affine.dma_start %{{.*}}[
   // CHECK:    affine.dma_wait [[TAG_ARG0_NESTED]]
   // CHECK:    affine.dma_wait [[TAG_ARG1_NESTED]]
   // CHECK:    affine.for %{{.*}} = 0 to 4 {
   // CHECK:      "foo"() : () -> ()
   // CHECK:  affine.dma_wait [[TAG_ARG0_NESTED]]
   // CHECK:  affine.dma_wait [[TAG_ARG1_NESTED]]
   // CHECK:  affine.for %{{.*}} = 0 to 4 {
   }
   memref.dealloc %5 : memref<2xi32>
   memref.dealloc %4 : memref<2xi32>
   memref.dealloc %3 : memref<2xi32>
   memref.dealloc %2 : memref<64x4xvector<8xf32>, 2>
   memref.dealloc %1 : memref<64x4xvector<8xf32>, 2>
   memref.dealloc %0 : memref<64x4xvector<8xf32>, 2>
   return
 // CHECK: }
 // CHECK-DAG:  memref.dealloc [[TAG_ARG1_NESTED]] : memref<2x2xi32>
 // CHECK-DAG:  memref.dealloc [[TAG_ARG0_NESTED]] : memref<2x2xi32>
 // CHECK-DAG:  memref.dealloc [[BUF_ARG1_NESTED]] : memref<2x64x4xvector<8xf32>, 2>
 // CHECK-DAG:  memref.dealloc [[BUF_ARG0_NESTED]] : memref<2x64x4xvector<8xf32>, 2>
 // CHECK-DAG:  memref.dealloc [[TAG_ARG2]] : memref<2x2xi32>
 // CHECK-DAG:  memref.dealloc [[BUF_ARG2]] : memref<2x64x4xvector<8xf32>, 2>
 // CHECK-NEXT: return
 }
 
 // -----
 #map2 = affine_map<(d0) -> ((d0 * 2048) floordiv 32)>
 
 // CHECK: func @loop_dma_dependent
 func.func @loop_dma_dependent(%arg2: memref<512x32xvector<8xf32>>) {
   %num_elts = arith.constant 256 : index
   %c0 = arith.constant 0 : index
   %0 = memref.alloc() : memref<64x4xvector<8xf32>, 2>
   %1 = memref.alloc() : memref<64x4xvector<8xf32>, 2>
   %2 = memref.alloc() : memref<64x4xvector<8xf32>, 2>
   %3 = memref.alloc() : memref<2xi32>
   %4 = memref.alloc() : memref<2xi32>
   %5 = memref.alloc() : memref<2xi32>
 
   // The two DMAs below are dependent (incoming and outgoing on the same
   // memref) in the same iteration; so no pipelining here.
   // CHECK-NOT: affine.dma_start
   // CHECK: affine.for %{{.*}} = 0 to 8 {
   affine.for %i0 = 0 to 8 {
     %6 = affine.apply #map2(%i0)
     affine.dma_start %arg2[%6, %c0], %2[%c0, %c0], %5[%c0], %num_elts : memref<512x32xvector<8xf32>>, memref<64x4xvector<8xf32>, 2>, memref<2xi32>
     affine.dma_wait %5[%c0], %num_elts : memref<2xi32>
 
     affine.dma_start %2[%c0, %c0], %arg2[%6, %c0], %5[%c0], %num_elts : memref<64x4xvector<8xf32>, 2>, memref<512x32xvector<8xf32>>, memref<2xi32>
     affine.dma_wait %5[%c0], %num_elts : memref<2xi32>
   }
   memref.dealloc %5 : memref<2xi32>
   memref.dealloc %4 : memref<2xi32>
   memref.dealloc %3 : memref<2xi32>
   memref.dealloc %2 : memref<64x4xvector<8xf32>, 2>
   memref.dealloc %1 : memref<64x4xvector<8xf32>, 2>
   memref.dealloc %0 : memref<64x4xvector<8xf32>, 2>
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @escaping_use
 func.func @escaping_use(%arg0: memref<512 x 32 x f32>) {
   %c32 = arith.constant 32 : index
   %num_elt = arith.constant 512 : index
   %zero = arith.constant 0 : index
   %Av = memref.alloc() : memref<32 x 32 x f32, 2>
   %tag = memref.alloc() : memref<1 x i32>
 
   // CHECK-NOT: affine.dma_start
   // CHECK: affine.for %{{.*}} = 0 to 16 {
   affine.for %kTT = 0 to 16 {
     affine.dma_start %arg0[%zero, %zero], %Av[%zero, %zero], %tag[%zero], %num_elt :
       memref<512 x 32 x f32>,
       memref<32 x 32 x f32, 2>, memref<1 x i32>
     affine.dma_wait %tag[%zero], %num_elt : memref<1 x i32>
     // escaping use; no DMA pipelining / double buffering will be done.
     "foo"(%Av) : (memref<32 x 32 x f32, 2>) -> ()
   }
   memref.dealloc %tag : memref<1 x i32>
   memref.dealloc %Av : memref<32 x 32 x f32, 2>
   return
 // CHECK:        "foo"(%{{[0-9a-zA-Z_]+}}) : (memref<32x32xf32, 2>) -> ()
 // CHECK:      }
 // CHECK:      return
 }
 
 // -----
 
 // CHECK-LABEL: func @escaping_tag
 func.func @escaping_tag(%arg0: memref<512 x 32 x f32>) {
   %c32 = arith.constant 32 : index
   %num_elt = arith.constant 512 : index
   %zero = arith.constant 0 : index
   %Av = memref.alloc() : memref<32 x 32 x f32, 2>
   %tag = memref.alloc() : memref<1 x i32>
 
   // CHECK-NOT: affine.dma_start
   // CHECK: affine.for %{{.*}} = 0 to 16 {
   affine.for %kTT = 0 to 16 {
     affine.dma_start %arg0[%zero, %zero], %Av[%zero, %zero], %tag[%zero], %num_elt :
       memref<512 x 32 x f32>,
       memref<32 x 32 x f32, 2>, memref<1 x i32>
     affine.dma_wait %tag[%zero], %num_elt : memref<1 x i32>
     // escaping use; no DMA pipelining / double buffering will be done.
     "foo"(%tag) : (memref<1 x i32>) -> ()
   }
   memref.dealloc %tag : memref<1 x i32>
   memref.dealloc %Av : memref<32 x 32 x f32, 2>
   return
 // CHECK:        "foo"(%{{[0-9a-zA-Z_]+}}) : (memref<1xi32>) -> ()
 // CHECK:      }
 // CHECK:      return
 }
 
 
 // -----
 
 // CHECK-LABEL: func @live_out_use
 func.func @live_out_use(%arg0: memref<512 x 32 x f32>) -> f32 {
   %c32 = arith.constant 32 : index
   %num_elt = arith.constant 512 : index
   %zero = arith.constant 0 : index
   %Av = memref.alloc() : memref<32 x 32 x f32, 2>
   %tag = memref.alloc() : memref<1 x i32>
 
   // CHECK-NOT: affine.dma_start
   // CHECK: affine.for %{{.*}} = 0 to 16 {
   affine.for %kTT = 0 to 16 {
     affine.dma_start %arg0[%zero, %zero], %Av[%zero, %zero], %tag[%zero], %num_elt :
       memref<512 x 32 x f32>,
       memref<32 x 32 x f32, 2>, memref<1 x i32>
     affine.dma_wait %tag[%zero], %num_elt : memref<1 x i32>
   }
   // Use live out of 'affine.for' op; no DMA pipelining will be done.
   %v = affine.load %Av[%zero, %zero] : memref<32 x 32 x f32, 2>
   memref.dealloc %tag : memref<1 x i32>
   memref.dealloc %Av : memref<32 x 32 x f32, 2>
   return %v : f32
 // CHECK:      affine.load %{{[0-9a-zA-Z_]+}}[%{{.*}}, %{{.*}}] : memref<32x32xf32, 2>
 // CHECK:      return
 }
 
 // -----
 
 // CHECK-LABEL: func @dynamic_shape_dma_buffer
 func.func @dynamic_shape_dma_buffer(%arg0: memref<512 x 32 x f32>, %Av: memref<? x ? x f32, 2>) {
   %num_elt = arith.constant 512 : index
   %zero = arith.constant 0 : index
   %tag = memref.alloc() : memref<1 x i32>
 
 // Double buffering for dynamic shaped buffer.
 // Note: Cannot capture C0 because there are multiple C0 constants in the IR.
 // CHECK:       memref.dim %{{.*}}, %{{.*}} : memref<?x?xf32, 2>
 // CHECK-NEXT:  %[[C1:.*]] = arith.constant 1 : index
 // CHECK-NEXT:  memref.dim %{{.*}}, %[[C1]] : memref<?x?xf32, 2>
 // CHECK-NEXT:  memref.alloc(%{{.*}}, %{{.*}}) : memref<2x?x?xf32, 2>
 // CHECK:       affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}} mod 2, 0, 0], %{{.*}}[%{{.*}} mod 2, 0], %{{.*}}
   affine.for %kTT = 0 to 16 {
     affine.dma_start %arg0[%zero, %zero], %Av[%zero, %zero], %tag[%zero], %num_elt :
       memref<512 x 32 x f32>,
       memref<? x ? x f32, 2>, memref<1 x i32>
     affine.dma_wait %tag[%zero], %num_elt : memref<1 x i32>
   }
   return
 // CHECK-NEXT:  affine.for %{{.*}} = 1 to 16 {
 // CHECK:         affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}} mod 2, 0, 0], %{{.*}}[%{{.*}} mod 2, 0], %{{.*}}
 // CHECK:         affine.dma_wait %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<2x1xi32>
 // CHECK:       }
 // CHECK:       affine.dma_wait %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<2x1xi32>
 // CHECK:       return
 }
 
 // Memref replacement will fail here due to a non-dereferencing use. However,
 // no incorrect transformation is performed in spite of one of the uses being a
 // dereferencing one since replaceAllMemRefUsesWith checks for escaping uses
 // before performing any replacement.
 // CHECK-LABEL: func @escaping_and_indexed_use_mix
 func.func @escaping_and_indexed_use_mix() {
   %A = memref.alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 0>
   %Ah = memref.alloc() : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
   %tag = memref.alloc() : memref<1 x f32>
   %zero = arith.constant 0 : index
   %num_elts = arith.constant 32 : index
 
   // alloc for the buffer is created but no replacement should happen.
   affine.for %i = 0 to 8 {
     affine.dma_start %A[%i], %Ah[%i], %tag[%zero], %num_elts : memref<256 x f32>, memref<32 x f32, 1>, memref<1 x f32>
     affine.dma_wait %tag[%zero], %num_elts : memref<1 x f32>
     "compute"(%Ah) : (memref<32 x f32, 1>) -> ()
     %v = affine.load %Ah[%i] : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
     "foo"(%v) : (f32) -> ()
   }
   memref.dealloc %A : memref<256 x f32, affine_map<(d0) -> (d0)>, 0>
   memref.dealloc %Ah : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
   return
 }
 // No replacement.
 // CHECK: affine.for %{{.*}} = 0 to 8 {
 // CHECK-NEXT:   affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}
 // CHECK-NEXT:   affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xf32>
 // CHECK-NEXT:   "compute"(%{{.*}}) : (memref<32xf32, 1>) -> ()
 // CHECK-NEXT:   [[VAL:%[0-9a-zA-Z_]+]] = affine.load %{{.*}}[%{{.*}}] : memref<32xf32, 1>
 // CHECK-NEXT:   "foo"([[VAL]]) : (f32) -> ()

mlir/test/Dialect/MemRef/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.func @test_norm_dynamic12(%arg0 : memref<1x?x?x14xf32, #map_tiled>) -> () {
     %c1 = arith.constant 1 : index
     %c2 = arith.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_:%.+]] = 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?x1x?x64xf32>
     // CHECK-DAG:       [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x1x?x64xf32>
     // CHECK-DAG:       [[CST_1_1_:%.+]] = arith.constant 1 : index
     // CHECK-DAG:       [[CST_14_:%.+]] = arith.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.func @test_norm_dynamic1234(%arg0 : memref<?x?x?x?xf32, #map_tiled1>) -> () {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
     %c2 = arith.constant 2 : index
     %c3 = arith.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_:%.+]] = arith.constant 0 : index
     // CHECK-DAG:       [[CST_1_:%.+]] = arith.constant 1 : index
     // CHECK-DAG:       [[CST_2_:%.+]] = arith.constant 2 : index
     // CHECK-DAG:       [[CST_3_:%.+]] = arith.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.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:           [[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.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:           [[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.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:           [[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.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:           [[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.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:           [[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
 }

mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir

 // RUN: mlir-opt -normalize-memrefs %s | FileCheck %s
 
 // For all these cases, we test if MemRefs Normalization works with the test
 // operations.
 // * test.op_norm: this operation has the MemRefsNormalizable attribute. The tests
 //   that include this operation are constructed so that the normalization should
 //   happen.
 // * test_op_nonnorm: this operation does not have the MemRefsNormalization
 //   attribute. The tests that include this operation are constructed so that the
 //    normalization should not happen.
 
 #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 64, d2 mod 32, d3 mod 64)>
 
 // Test with op_norm and maps in arguments and in the operations in the function.
 
 // CHECK-LABEL: test_norm
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>)
 func.func @test_norm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
     %0 = memref.alloc() : memref<1x16x14x14xf32, #map0>
     "test.op_norm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
     memref.dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
     // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x1x1x32x64xf32>
     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
     // CHECK: memref.dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
     return
 }
 
 // Same test with op_nonnorm, with maps in the arguments and the operations in the function.
 
 // CHECK-LABEL: test_nonnorm
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32, #[[MAP:.*]]>)
 func.func @test_nonnorm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
     %0 = memref.alloc() : memref<1x16x14x14xf32, #map0>
     "test.op_nonnorm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
     memref.dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
     // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x14x14xf32, #[[MAP]]>
     // CHECK: "test.op_nonnorm"(%[[ARG0]], %[[v0]]) : (memref<1x16x14x14xf32, #[[MAP]]>, memref<1x16x14x14xf32, #[[MAP]]>) -> ()
     // CHECK: memref.dealloc %[[v0]] : memref<1x16x14x14xf32, #[[MAP]]>
     return
 }
 
 // Test with op_nonnorm whose memref map layouts are identity. This op_nonnorm
 // does not block the normalization of other operations.
 
 // CHECK-LABEL: test_nonnorm_identity_layout
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>)
 func.func @test_nonnorm_identity_layout(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
     %0 = memref.alloc() : memref<1x16x14x14xf32>
     "test.op_nonnorm"(%0, %0) : (memref<1x16x14x14xf32>, memref<1x16x14x14xf32>) -> ()
     "test.op_norm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32>) -> ()
     memref.dealloc %0 :  memref<1x16x14x14xf32>
 
     // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x14x14xf32>
     // CHECK: "test.op_nonnorm"(%[[v0]], %[[v0]]) : (memref<1x16x14x14xf32>, memref<1x16x14x14xf32>) -> ()
     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x14x14xf32>) -> ()
     // CHECK: memref.dealloc %[[v0]] : memref<1x16x14x14xf32>
     return
 }
 
 // Test with op_norm, with maps in the operations in the function.
 
 // CHECK-LABEL: test_norm_mix
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>
 func.func @test_norm_mix(%arg0 : memref<1x16x1x1x32x64xf32>) -> () {
     %0 = memref.alloc() : memref<1x16x14x14xf32, #map0>
     "test.op_norm"(%arg0, %0) : (memref<1x16x1x1x32x64xf32>, memref<1x16x14x14xf32, #map0>) -> ()
     memref.dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
     // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x1x1x32x64xf32>
     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
     // CHECK: memref.dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
     return
 }
 
 // Test with maps in load and store ops.
 
 #map_tile = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 32, d2 mod 32, d3 mod 32)>
 
 // CHECK-LABEL: test_load_store
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32>
 func.func @test_load_store(%arg0 : memref<1x16x14x14xf32>) -> () {
     %0 = memref.alloc() : memref<1x16x14x14xf32, #map_tile>
     // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x1x1x32x32xf32>
     %1 = memref.alloc() : memref<1x16x14x14xf32>
     // CHECK: %[[v1:.*]] = memref.alloc() : memref<1x16x14x14xf32>
     "test.op_norm"(%0, %1) : (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>) -> ()
     // CHECK: "test.op_norm"(%[[v0]], %[[v1]]) : (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>) -> ()
     %cst = arith.constant 3.0 : f32
     affine.for %i = 0 to 1 {
       affine.for %j = 0 to 16 {
         affine.for %k = 0 to 14 {
           affine.for %l = 0 to 14 {
             %2 = memref.load %1[%i, %j, %k, %l] : memref<1x16x14x14xf32>
             // CHECK: memref<1x16x14x14xf32>
             %3 = arith.addf %2, %cst : f32
             memref.store %3, %arg0[%i, %j, %k, %l] : memref<1x16x14x14xf32>
             // CHECK: memref<1x16x14x14xf32>
           }
         }
       }
     }
     memref.dealloc %0 :  memref<1x16x14x14xf32, #map_tile>
     // CHECK: memref.dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
     memref.dealloc %1 :  memref<1x16x14x14xf32>
     // CHECK: memref.dealloc %[[v1]] : memref<1x16x14x14xf32>
     return
 }
 
 // Test with op_norm_ret, with maps in the results of normalizable operation.
 
 // CHECK-LABEL: test_norm_ret
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x32xf32>) -> (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>) {
 func.func @test_norm_ret(%arg0: memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>) {
     %0 = memref.alloc() : memref<1x16x14x14xf32, #map_tile>
     // CHECK-NEXT: %[[v0:.*]] = memref.alloc() : memref<1x16x1x1x32x32xf32>
     %1, %2 = "test.op_norm_ret"(%arg0) : (memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>)
     // CHECK-NEXT: %[[v1:.*]], %[[v2:.*]] = "test.op_norm_ret"
     // CHECK-SAME: (memref<1x16x1x1x32x32xf32>) -> (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>)
     "test.op_norm"(%1, %0) : (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32, #map_tile>) -> ()
     // CHECK-NEXT: "test.op_norm"
     // CHECK-SAME: : (memref<1x16x1x1x32x32xf32>, memref<1x16x1x1x32x32xf32>) -> ()
     memref.dealloc %0 : memref<1x16x14x14xf32, #map_tile>
     // CHECK-NEXT: memref.dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
     return %1, %2 : memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>
     // CHECK-NEXT: return %[[v1]], %[[v2]] : memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>
 }
 
 // Test with an arbitrary op that references the function symbol.
 
 "test.op_funcref"() {func = @test_norm_mix} : () -> ()
 
 
 // -----
 
 #map_1d_tile = affine_map<(d0) -> (d0 floordiv 32, d0 mod 32)>
 
 // Test with memref.reinterpret_cast
 
 // CHECK-LABEL: test_norm_reinterpret_cast
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x32xf32>) -> memref<3x1x1xf32> {
 func.func @test_norm_reinterpret_cast(%arg0 : memref<3xf32, #map_1d_tile>) -> (memref<3x1x1xf32>) {
     %0 = memref.alloc() : memref<3xf32>
     "test.op_norm"(%arg0, %0) : (memref<3xf32, #map_1d_tile>, memref<3xf32>) -> ()
     %1 = memref.reinterpret_cast %0 to offset: [0], sizes: [3, 1, 1], strides: [1, 1, 1] : memref<3xf32> to memref<3x1x1xf32>
     // CHECK: %[[v0:.*]] = memref.alloc() : memref<3xf32>
     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x32xf32>, memref<3xf32>) -> ()
     // CHECK: memref.reinterpret_cast %[[v0]] to offset: [0], sizes: [3, 1, 1], strides: [1, 1, 1] : memref<3xf32> to memref<3x1x1xf32>
     return %1 : memref<3x1x1xf32>
 }

mlir/test/Dialect/MemRef/normalize-memrefs.mlir

 // RUN: mlir-opt -normalize-memrefs -allow-unregistered-dialect %s | FileCheck %s
 
 // This file tests whether the memref type having non-trivial map layouts
 // are normalized to trivial (identity) layouts.
 
 // CHECK-LABEL: func @permute()
 func.func @permute() {
   %A = memref.alloc() : memref<64x256xf32, affine_map<(d0, d1) -> (d1, d0)>>
   affine.for %i = 0 to 64 {
     affine.for %j = 0 to 256 {
       %1 = affine.load %A[%i, %j] : memref<64x256xf32, affine_map<(d0, d1) -> (d1, d0)>>
       "prevent.dce"(%1) : (f32) -> ()
     }
   }
   memref.dealloc %A : memref<64x256xf32, affine_map<(d0, d1) -> (d1, d0)>>
   return
 }
 // The old memref alloc should disappear.
 // CHECK-NOT:  memref<64x256xf32>
 // CHECK:      [[MEM:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<256x64xf32>
 // CHECK-NEXT: affine.for %[[I:arg[0-9a-zA-Z_]+]] = 0 to 64 {
 // CHECK-NEXT:   affine.for %[[J:arg[0-9a-zA-Z_]+]] = 0 to 256 {
 // CHECK-NEXT:     affine.load [[MEM]][%[[J]], %[[I]]] : memref<256x64xf32>
 // CHECK-NEXT:     "prevent.dce"
 // CHECK-NEXT:   }
 // CHECK-NEXT: }
 // CHECK-NEXT: memref.dealloc [[MEM]]
 // CHECK-NEXT: return
 
 // CHECK-LABEL: func @shift
 func.func @shift(%idx : index) {
   // CHECK-NEXT: memref.alloc() : memref<65xf32>
   %A = memref.alloc() : memref<64xf32, affine_map<(d0) -> (d0 + 1)>>
   // CHECK-NEXT: affine.load %{{.*}}[symbol(%arg0) + 1] : memref<65xf32>
   affine.load %A[%idx] : memref<64xf32, affine_map<(d0) -> (d0 + 1)>>
   affine.for %i = 0 to 64 {
     %1 = affine.load %A[%i] : memref<64xf32, affine_map<(d0) -> (d0 + 1)>>
     "prevent.dce"(%1) : (f32) -> ()
     // CHECK: %{{.*}} = affine.load %{{.*}}[%arg{{.*}} + 1] : memref<65xf32>
   }
   return
 }
 
 // CHECK-LABEL: func @high_dim_permute()
 func.func @high_dim_permute() {
   // CHECK-NOT: memref<64x128x256xf32,
   %A = memref.alloc() : memref<64x128x256xf32, affine_map<(d0, d1, d2) -> (d2, d0, d1)>>
   // CHECK: %[[I:arg[0-9a-zA-Z_]+]]
   affine.for %i = 0 to 64 {
     // CHECK: %[[J:arg[0-9a-zA-Z_]+]]
     affine.for %j = 0 to 128 {
       // CHECK: %[[K:arg[0-9a-zA-Z_]+]]
       affine.for %k = 0 to 256 {
         %1 = affine.load %A[%i, %j, %k] : memref<64x128x256xf32, affine_map<(d0, d1, d2) -> (d2, d0, d1)>>
         // CHECK: %{{.*}} = affine.load %{{.*}}[%[[K]], %[[I]], %[[J]]] : memref<256x64x128xf32>
         "prevent.dce"(%1) : (f32) -> ()
       }
     }
   }
   return
 }
 
 // CHECK-LABEL: func @invalid_map
 func.func @invalid_map() {
   %A = memref.alloc() : memref<64x128xf32, affine_map<(d0, d1) -> (d0, -d1 - 10)>>
   // CHECK: %{{.*}} = memref.alloc() : memref<64x128xf32,
   return
 }
 
 // A tiled layout.
 // CHECK-LABEL: func @data_tiling
 func.func @data_tiling(%idx : index) {
   // CHECK: memref.alloc() : memref<8x32x8x16xf32>
   %A = memref.alloc() : memref<64x512xf32, affine_map<(d0, d1) -> (d0 floordiv 8, d1 floordiv 16, d0 mod 8, d1 mod 16)>>
   // CHECK: affine.load %{{.*}}[symbol(%arg0) floordiv 8, symbol(%arg0) floordiv 16, symbol(%arg0) mod 8, symbol(%arg0) mod 16]
   %1 = affine.load %A[%idx, %idx] : memref<64x512xf32, affine_map<(d0, d1) -> (d0 floordiv 8, d1 floordiv 16, d0 mod 8, d1 mod 16)>>
   "prevent.dce"(%1) : (f32) -> ()
   return
 }
 
 // Strides 2 and 4 along respective dimensions.
 // CHECK-LABEL: func @strided
 func.func @strided() {
   %A = memref.alloc() : memref<64x128xf32, affine_map<(d0, d1) -> (2*d0, 4*d1)>>
   // CHECK: affine.for %[[IV0:.*]] =
   affine.for %i = 0 to 64 {
     // CHECK: affine.for %[[IV1:.*]] =
     affine.for %j = 0 to 128 {
       // CHECK: affine.load %{{.*}}[%[[IV0]] * 2, %[[IV1]] * 4] : memref<127x509xf32>
       %1 = affine.load %A[%i, %j] : memref<64x128xf32, affine_map<(d0, d1) -> (2*d0, 4*d1)>>
       "prevent.dce"(%1) : (f32) -> ()
     }
   }
   return
 }
 
 // Strided, but the strides are in the linearized space.
 // CHECK-LABEL: func @strided_cumulative
 func.func @strided_cumulative() {
   %A = memref.alloc() : memref<2x5xf32, affine_map<(d0, d1) -> (3*d0 + 17*d1)>>
   // CHECK: affine.for %[[IV0:.*]] =
   affine.for %i = 0 to 2 {
     // CHECK: affine.for %[[IV1:.*]] =
     affine.for %j = 0 to 5 {
       // CHECK: affine.load %{{.*}}[%[[IV0]] * 3 + %[[IV1]] * 17] : memref<72xf32>
       %1 = affine.load %A[%i, %j]  : memref<2x5xf32, affine_map<(d0, d1) -> (3*d0 + 17*d1)>>
       "prevent.dce"(%1) : (f32) -> ()
     }
   }
   return
 }
 
 // Symbolic operand for alloc, although unused. Tests replaceAllMemRefUsesWith
 // when the index remap has symbols.
 // CHECK-LABEL: func @symbolic_operands
 func.func @symbolic_operands(%s : index) {
   // CHECK: memref.alloc() : memref<100xf32>
   %A = memref.alloc()[%s] : memref<10x10xf32, affine_map<(d0,d1)[s0] -> (10*d0 + d1)>>
   affine.for %i = 0 to 10 {
     affine.for %j = 0 to 10 {
       // CHECK: affine.load %{{.*}}[%{{.*}} * 10 + %{{.*}}] : memref<100xf32>
       %1 = affine.load %A[%i, %j] : memref<10x10xf32, affine_map<(d0,d1)[s0] -> (10*d0 + d1)>>
       "prevent.dce"(%1) : (f32) -> ()
     }
   }
   return
 }
 
 // Semi-affine maps, normalization not implemented yet.
 // CHECK-LABEL: func @semi_affine_layout_map
 func.func @semi_affine_layout_map(%s0: index, %s1: index) {
   %A = memref.alloc()[%s0, %s1] : memref<256x1024xf32, affine_map<(d0, d1)[s0, s1] -> (d0*s0 + d1*s1)>>
   affine.for %i = 0 to 256 {
     affine.for %j = 0 to 1024 {
       // CHECK: memref<256x1024xf32, #map{{[0-9a-zA-Z_]+}}>
       affine.load %A[%i, %j] : memref<256x1024xf32, affine_map<(d0, d1)[s0, s1] -> (d0*s0 + d1*s1)>>
     }
   }
   return
 }
 
 // CHECK-LABEL: func @alignment
 func.func @alignment() {
   %A = memref.alloc() {alignment = 32 : i64}: memref<64x128x256xf32, affine_map<(d0, d1, d2) -> (d2, d0, d1)>>
   // CHECK-NEXT: memref.alloc() {alignment = 32 : i64} : memref<256x64x128xf32>
   return
 }
 
 #tile = affine_map < (i)->(i floordiv 4, i mod 4) >
 
 // Following test cases check the inter-procedural memref normalization.
 
 // Test case 1: Check normalization for multiple memrefs in a function argument list.
 // CHECK-LABEL: func @multiple_argument_type
 // CHECK-SAME:  (%[[A:arg[0-9a-zA-Z_]+]]: memref<4x4xf64>, %[[B:arg[0-9a-zA-Z_]+]]: f64, %[[C:arg[0-9a-zA-Z_]+]]: memref<2x4xf64>, %[[D:arg[0-9a-zA-Z_]+]]: memref<24xf64>) -> f64
 func.func @multiple_argument_type(%A: memref<16xf64, #tile>, %B: f64, %C: memref<8xf64, #tile>, %D: memref<24xf64>) -> f64 {
   %a = affine.load %A[0] : memref<16xf64, #tile>
   %p = arith.mulf %a, %a : f64
   affine.store %p, %A[10] : memref<16xf64, #tile>
   call @single_argument_type(%C): (memref<8xf64, #tile>) -> ()
   return %B : f64
 }
 
 // CHECK: %[[a:[0-9a-zA-Z_]+]] = affine.load %[[A]][0, 0] : memref<4x4xf64>
 // CHECK: %[[p:[0-9a-zA-Z_]+]] = arith.mulf %[[a]], %[[a]] : f64
 // CHECK: affine.store %[[p]], %[[A]][2, 2] : memref<4x4xf64>
 // CHECK: call @single_argument_type(%[[C]]) : (memref<2x4xf64>) -> ()
 // CHECK: return %[[B]] : f64
 
 // Test case 2: Check normalization for single memref argument in a function.
 // CHECK-LABEL: func @single_argument_type
 // CHECK-SAME: (%[[C:arg[0-9a-zA-Z_]+]]: memref<2x4xf64>)
 func.func @single_argument_type(%C : memref<8xf64, #tile>) {
   %a = memref.alloc(): memref<8xf64, #tile>
   %b = memref.alloc(): memref<16xf64, #tile>
   %d = arith.constant 23.0 : f64
   %e = memref.alloc(): memref<24xf64>
   call @single_argument_type(%a): (memref<8xf64, #tile>) -> ()
   call @single_argument_type(%C): (memref<8xf64, #tile>) -> ()
   call @multiple_argument_type(%b, %d, %a, %e): (memref<16xf64, #tile>, f64, memref<8xf64, #tile>, memref<24xf64>) -> f64
   return
 }
 
 // CHECK: %[[a:[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x4xf64>
 // CHECK: %[[b:[0-9a-zA-Z_]+]] = memref.alloc() : memref<4x4xf64>
 // CHECK: %cst = arith.constant 2.300000e+01 : f64
 // CHECK: %[[e:[0-9a-zA-Z_]+]] = memref.alloc() : memref<24xf64>
 // CHECK: call @single_argument_type(%[[a]]) : (memref<2x4xf64>) -> ()
 // CHECK: call @single_argument_type(%[[C]]) : (memref<2x4xf64>) -> ()
 // CHECK: call @multiple_argument_type(%[[b]], %cst, %[[a]], %[[e]]) : (memref<4x4xf64>, f64, memref<2x4xf64>, memref<24xf64>) -> f64
 
 // Test case 3: Check function returning any other type except memref.
 // CHECK-LABEL: func @non_memref_ret
 // CHECK-SAME: (%[[C:arg[0-9a-zA-Z_]+]]: memref<2x4xf64>) -> i1
 func.func @non_memref_ret(%A: memref<8xf64, #tile>) -> i1 {
   %d = arith.constant 1 : i1
   return %d : i1
 }
 
 // Test cases here onwards deal with normalization of memref in function signature, caller site.
 
 // Test case 4: Check successful memref normalization in case of inter/intra-recursive calls.
 // CHECK-LABEL: func @ret_multiple_argument_type
 // CHECK-SAME: (%[[A:arg[0-9a-zA-Z_]+]]: memref<4x4xf64>, %[[B:arg[0-9a-zA-Z_]+]]: f64, %[[C:arg[0-9a-zA-Z_]+]]: memref<2x4xf64>) -> (memref<2x4xf64>, f64)
 func.func @ret_multiple_argument_type(%A: memref<16xf64, #tile>, %B: f64, %C: memref<8xf64, #tile>) -> (memref<8xf64, #tile>, f64) {
   %a = affine.load %A[0] : memref<16xf64, #tile>
   %p = arith.mulf %a, %a : f64
   %cond = arith.constant 1 : i1
   cf.cond_br %cond, ^bb1, ^bb2
   ^bb1:
     %res1, %res2 = call @ret_single_argument_type(%C) : (memref<8xf64, #tile>) -> (memref<16xf64, #tile>, memref<8xf64, #tile>)
     return %res2, %p: memref<8xf64, #tile>, f64
   ^bb2:
     return %C, %p: memref<8xf64, #tile>, f64
 }
 
 // CHECK:   %[[a:[0-9a-zA-Z_]+]] = affine.load %[[A]][0, 0] : memref<4x4xf64>
 // CHECK:   %[[p:[0-9a-zA-Z_]+]] = arith.mulf %[[a]], %[[a]] : f64
 // CHECK:   %true = arith.constant true
 // CHECK:   cf.cond_br %true, ^bb1, ^bb2
 // CHECK: ^bb1:  // pred: ^bb0
 // CHECK:   %[[res:[0-9a-zA-Z_]+]]:2 = call @ret_single_argument_type(%[[C]]) : (memref<2x4xf64>) -> (memref<4x4xf64>, memref<2x4xf64>)
 // CHECK:   return %[[res]]#1, %[[p]] : memref<2x4xf64>, f64
 // CHECK: ^bb2:  // pred: ^bb0
 // CHECK:   return %{{.*}}, %{{.*}} : memref<2x4xf64>, f64
 
 // CHECK-LABEL: func @ret_single_argument_type
 // CHECK-SAME: (%[[C:arg[0-9a-zA-Z_]+]]: memref<2x4xf64>) -> (memref<4x4xf64>, memref<2x4xf64>)
 func.func @ret_single_argument_type(%C: memref<8xf64, #tile>) -> (memref<16xf64, #tile>, memref<8xf64, #tile>){
   %a = memref.alloc() : memref<8xf64, #tile>
   %b = memref.alloc() : memref<16xf64, #tile>
   %d = arith.constant 23.0 : f64
   call @ret_single_argument_type(%a) : (memref<8xf64, #tile>) -> (memref<16xf64, #tile>, memref<8xf64, #tile>)
   call @ret_single_argument_type(%C) : (memref<8xf64, #tile>) -> (memref<16xf64, #tile>, memref<8xf64, #tile>)
   %res1, %res2 = call @ret_multiple_argument_type(%b, %d, %a) : (memref<16xf64, #tile>, f64, memref<8xf64, #tile>) -> (memref<8xf64, #tile>, f64)
   %res3, %res4 = call @ret_single_argument_type(%res1) : (memref<8xf64, #tile>) -> (memref<16xf64, #tile>, memref<8xf64, #tile>)
   return %b, %a: memref<16xf64, #tile>, memref<8xf64, #tile>
 }
 
 // CHECK: %[[a:[0-9a-zA-Z_]+]] = memref.alloc() : memref<2x4xf64>
 // CHECK: %[[b:[0-9a-zA-Z_]+]] = memref.alloc() : memref<4x4xf64>
 // CHECK: %cst = arith.constant 2.300000e+01 : f64
 // CHECK: %[[resA:[0-9a-zA-Z_]+]]:2 = call @ret_single_argument_type(%[[a]]) : (memref<2x4xf64>) -> (memref<4x4xf64>, memref<2x4xf64>)
 // CHECK: %[[resB:[0-9a-zA-Z_]+]]:2 = call @ret_single_argument_type(%[[C]]) : (memref<2x4xf64>) -> (memref<4x4xf64>, memref<2x4xf64>)
 // CHECK: %[[resC:[0-9a-zA-Z_]+]]:2 = call @ret_multiple_argument_type(%[[b]], %cst, %[[a]]) : (memref<4x4xf64>, f64, memref<2x4xf64>) -> (memref<2x4xf64>, f64)
 // CHECK: %[[resD:[0-9a-zA-Z_]+]]:2 = call @ret_single_argument_type(%[[resC]]#0) : (memref<2x4xf64>) -> (memref<4x4xf64>, memref<2x4xf64>)
 // CHECK: return %{{.*}}, %{{.*}} : memref<4x4xf64>, memref<2x4xf64>
 
 // Test case set #5: To check normalization in a chain of interconnected functions.
 // CHECK-LABEL: func @func_A
 // CHECK-SAME: (%[[A:arg[0-9a-zA-Z_]+]]: memref<2x4xf64>)
 func.func @func_A(%A: memref<8xf64, #tile>) {
   call @func_B(%A) : (memref<8xf64, #tile>) -> ()
   return
 }
 // CHECK: call @func_B(%[[A]]) : (memref<2x4xf64>) -> ()
 
 // CHECK-LABEL: func @func_B
 // CHECK-SAME: (%[[A:arg[0-9a-zA-Z_]+]]: memref<2x4xf64>)
 func.func @func_B(%A: memref<8xf64, #tile>) {
   call @func_C(%A) : (memref<8xf64, #tile>) -> ()
   return
 }
 // CHECK: call @func_C(%[[A]]) : (memref<2x4xf64>) -> ()
 
 // CHECK-LABEL: func @func_C
 // CHECK-SAME: (%[[A:arg[0-9a-zA-Z_]+]]: memref<2x4xf64>)
 func.func @func_C(%A: memref<8xf64, #tile>) {
   return
 }
 
 // Test case set #6: Checking if no normalization takes place in a scenario: A -> B -> C and B has an unsupported type.
 // CHECK-LABEL: func @some_func_A
 // CHECK-SAME: (%[[A:arg[0-9a-zA-Z_]+]]: memref<8xf64, #map{{[0-9a-zA-Z_]+}}>)
 func.func @some_func_A(%A: memref<8xf64, #tile>) {
   call @some_func_B(%A) : (memref<8xf64, #tile>) -> ()
   return
 }
 // CHECK: call @some_func_B(%[[A]]) : (memref<8xf64, #map{{[0-9a-zA-Z_]+}}>) -> ()
 
 // CHECK-LABEL: func @some_func_B
 // CHECK-SAME: (%[[A:arg[0-9a-zA-Z_]+]]: memref<8xf64, #map{{[0-9a-zA-Z_]+}}>)
 func.func @some_func_B(%A: memref<8xf64, #tile>) {
   "test.test"(%A) : (memref<8xf64, #tile>) -> ()
   call @some_func_C(%A) : (memref<8xf64, #tile>) -> ()
   return
 }
 // CHECK: call @some_func_C(%[[A]]) : (memref<8xf64, #map{{[0-9a-zA-Z_]+}}>) -> ()
 
 // CHECK-LABEL: func @some_func_C
 // CHECK-SAME: (%[[A:arg[0-9a-zA-Z_]+]]: memref<8xf64, #map{{[0-9a-zA-Z_]+}}>)
 func.func @some_func_C(%A: memref<8xf64, #tile>) {
   return
 }
 
 // Test case set #7: Check normalization in case of external functions.
 // CHECK-LABEL: func private @external_func_A
 // CHECK-SAME: (memref<4x4xf64>)
 func.func private @external_func_A(memref<16xf64, #tile>) -> ()
 
 // CHECK-LABEL: func private @external_func_B
 // CHECK-SAME: (memref<4x4xf64>, f64) -> memref<2x4xf64>
 func.func private @external_func_B(memref<16xf64, #tile>, f64) -> (memref<8xf64, #tile>)
 
 // CHECK-LABEL: func @simply_call_external()
 func.func @simply_call_external() {
   %a = memref.alloc() : memref<16xf64, #tile>
   call @external_func_A(%a) : (memref<16xf64, #tile>) -> ()
   return
 }
 // CHECK: %[[a:[0-9a-zA-Z_]+]] = memref.alloc() : memref<4x4xf64>
 // CHECK: call @external_func_A(%[[a]]) : (memref<4x4xf64>) -> ()
 
 // CHECK-LABEL: func @use_value_of_external
 // CHECK-SAME: (%[[A:arg[0-9a-zA-Z_]+]]: memref<4x4xf64>, %[[B:arg[0-9a-zA-Z_]+]]: f64) -> memref<2x4xf64>
 func.func @use_value_of_external(%A: memref<16xf64, #tile>, %B: f64) -> (memref<8xf64, #tile>) {
   %res = call @external_func_B(%A, %B) : (memref<16xf64, #tile>, f64) -> (memref<8xf64, #tile>)
   return %res : memref<8xf64, #tile>
 }
 // CHECK: %[[res:[0-9a-zA-Z_]+]] = call @external_func_B(%[[A]], %[[B]]) : (memref<4x4xf64>, f64) -> memref<2x4xf64>
 // CHECK: return %{{.*}} : memref<2x4xf64>
 
 // CHECK-LABEL: func @affine_parallel_norm
 func.func @affine_parallel_norm() ->  memref<8xf32, #tile> {
   %c = arith.constant 23.0 : f32
   %a = memref.alloc() : memref<8xf32, #tile>
   // CHECK: affine.parallel (%{{.*}}) = (0) to (8) reduce ("assign") -> (memref<2x4xf32>)
   %1 = affine.parallel (%i) = (0) to (8) reduce ("assign") ->  memref<8xf32, #tile> {
     affine.store %c, %a[%i] : memref<8xf32, #tile>
     // CHECK: affine.yield %{{.*}} : memref<2x4xf32>
     affine.yield %a : memref<8xf32, #tile>
   }
   return %1 : memref<8xf32, #tile>
 }
 
 #map = affine_map<(d0, d1)[s0] -> (d0 * 3 + s0 + d1)>
 // CHECK-LABEL: func.func @map_symbol
 func.func @map_symbol() -> memref<2x3xf32, #map> {
   %c1 = arith.constant 1 : index
   // The constant isn't propagated here and the utility can't compute a constant
   // upper bound for the memref dimension in the absence of that.
   // CHECK: memref.alloc()[%{{.*}}]
   %0 = memref.alloc()[%c1] : memref<2x3xf32, #map>
   return %0 : memref<2x3xf32, #map>
 }
 
 #neg = affine_map<(d0, d1) -> (d0, d1 - 100)>
 // CHECK-LABEL: func.func @neg_map
 func.func @neg_map() -> memref<2x3xf32, #neg> {
   // This isn't a valid map for normalization.
   // CHECK: memref.alloc() : memref<2x3xf32, #{{.*}}>
   %0 = memref.alloc() : memref<2x3xf32, #neg>
   return %0 : memref<2x3xf32, #neg>
 }
 
 // CHECK-LABEL: func @memref_with_strided_offset
 func.func @memref_with_strided_offset(%arg0: tensor<128x512xf32>, %arg1: index, %arg2: index) -> tensor<16x512xf32> {
   %c0 = arith.constant 0 : index
   %0 = bufferization.to_memref %arg0 : memref<128x512xf32, strided<[?, ?], offset: ?>>
   %subview = memref.subview %0[%arg2, 0] [%arg1, 512] [1, 1] : memref<128x512xf32, strided<[?, ?], offset: ?>> to memref<?x512xf32, strided<[?, ?], offset: ?>>
   // CHECK: %{{.*}} = memref.cast %{{.*}} : memref<?x512xf32, strided<[?, ?], offset: ?>> to memref<16x512xf32, strided<[?, ?], offset: ?>>
   %cast = memref.cast %subview : memref<?x512xf32, strided<[?, ?], offset: ?>> to memref<16x512xf32, strided<[?, ?], offset: ?>>
   %1 = bufferization.to_tensor %cast : memref<16x512xf32, strided<[?, ?], offset: ?>>
   return %1 : tensor<16x512xf32>
 }