// RUN: mlir-opt -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer fusion-maximal}))' %s | FileCheck %s

// Test fusion of affine nests inside other region-holding ops (scf.for in the
// test case below).

// CHECK-LABEL: func @fusion_inner_simple
func.func @fusion_inner_simple(%A : memref<10xf32>) {
  %cst = arith.constant 0.0 : f32

  affine.for %i = 0 to 100 {
    %B = memref.alloc() : memref<10xf32>
    %C = memref.alloc() : memref<10xf32>

    affine.for %j = 0 to 10 {
      %v = affine.load %A[%j] : memref<10xf32>
      affine.store %v, %B[%j] : memref<10xf32>
    }

    affine.for %j = 0 to 10 {
      %v = affine.load %B[%j] : memref<10xf32>
      affine.store %v, %C[%j] : memref<10xf32>
    }
  }

  // CHECK:      affine.for %{{.*}} = 0 to 100
  // CHECK-NEXT:   memref.alloc
  // CHECK-NEXT:   memref.alloc
  // CHECK-NEXT:   affine.for %{{.*}} = 0 to 10
  // CHECK-NOT:    affine.for

  return
}

// CHECK-LABEL: func @fusion_inner_simple_scf
func.func @fusion_inner_simple_scf(%A : memref<10xf32>) {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c100 = arith.constant 100 : index
  %cst = arith.constant 0.0 : f32

  scf.for %i = %c0 to %c100 step %c1 {
    %B = memref.alloc() : memref<10xf32>
    %C = memref.alloc() : memref<10xf32>

    affine.for %j = 0 to 10 {
      %v = affine.load %A[%j] : memref<10xf32>
      affine.store %v, %B[%j] : memref<10xf32>
    }

    affine.for %j = 0 to 10 {
      %v = affine.load %B[%j] : memref<10xf32>
      affine.store %v, %C[%j] : memref<10xf32>
    }
  }
  // CHECK:      scf.for
  // CHECK-NEXT:   memref.alloc
  // CHECK-NEXT:   memref.alloc
  // CHECK-NEXT:   affine.for %{{.*}} = 0 to 10
  // CHECK-NOT:    affine.for
  return
}

// CHECK-LABEL: func @fusion_inner_multiple_nests
func.func @fusion_inner_multiple_nests() {
  %alloc_5 = memref.alloc() {alignment = 64 : i64} : memref<4x4xi8>
  %alloc_10 = memref.alloc() : memref<8x4xi32>
  affine.for %arg8 = 0 to 4 {
    %alloc_14 = memref.alloc() : memref<4xi8>
    %alloc_15 = memref.alloc() : memref<8x4xi8>
    affine.for %arg9 = 0 to 4 {
      %0 = affine.load %alloc_5[%arg9, %arg8] : memref<4x4xi8>
      affine.store %0, %alloc_14[%arg9] : memref<4xi8>
    }
    %alloc_16 = memref.alloc() : memref<4xi8>
    affine.for %arg9 = 0 to 4 {
      %0 = affine.load %alloc_14[%arg9] : memref<4xi8>
      affine.store %0, %alloc_16[%arg9] : memref<4xi8>
    }
    affine.for %arg9 = 0 to 2 {
      %0 = affine.load %alloc_15[%arg9 * 4, 0] : memref<8x4xi8>
      %1 = affine.load %alloc_16[0] : memref<4xi8>
      %2 = affine.load %alloc_10[%arg9 * 4, %arg8] : memref<8x4xi32>
      %3 = arith.muli %0, %1 : i8
      %4 = arith.extsi %3 : i8 to i32
      %5 = arith.addi %4, %2 : i32
      affine.store %5, %alloc_10[%arg9 * 4 + 3, %arg8] : memref<8x4xi32>
    }
    memref.dealloc %alloc_16 : memref<4xi8>
  }
  // CHECK:      affine.for %{{.*}} = 0 to 4 {
  // Everything inside fused into two nests (the second will be DCE'd).
  // CHECK-NEXT:   memref.alloc() : memref<4xi8>
  // CHECK-NEXT:   memref.alloc() : memref<1xi8>
  // CHECK-NEXT:   memref.alloc() : memref<1xi8>
  // CHECK-NEXT:   memref.alloc() : memref<8x4xi8>
  // CHECK-NEXT:   memref.alloc() : memref<4xi8>
  // CHECK-NEXT:   affine.for %{{.*}} = 0 to 2 {
  // CHECK:        }
  // CHECK:        affine.for %{{.*}} = 0 to 4 {
  // CHECK:        }
  // CHECK-NEXT:   memref.dealloc
  // CHECK-NEXT: }
  // CHECK-NEXT: return
  return
}

// CHECK-LABEL: func @fusion_inside_scf_while
func.func @fusion_inside_scf_while(%A : memref<10xf32>) {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c100 = arith.constant 100 : index
  %cst = arith.constant 0.0 : f32

  %0 = scf.while (%arg3 = %cst) : (f32) -> (f32) {
    %1 = arith.cmpf ult, %arg3, %cst : f32
    scf.condition(%1) %arg3 : f32
  } do {
  ^bb0(%arg5: f32):

    %B = memref.alloc() : memref<10xf32>
    %C = memref.alloc() : memref<10xf32>

    affine.for %j = 0 to 10 {
      %v = affine.load %A[%j] : memref<10xf32>
      affine.store %v, %B[%j] : memref<10xf32>
    }

    affine.for %j = 0 to 10 {
      %v = affine.load %B[%j] : memref<10xf32>
      affine.store %v, %C[%j] : memref<10xf32>
    }
    %1 = arith.mulf %arg5, %cst : f32
    scf.yield %1 : f32
  }
  // CHECK:      scf.while
  // CHECK:        affine.for %{{.*}} = 0 to 10
  // CHECK-NOT:    affine.for
  // CHECK:        scf.yield
  return
}


memref.global "private" constant @__constant_10x2xf32 : memref<10x2xf32> = dense<0.000000e+00>

// CHECK-LABEL: func @fusion_inner_long
func.func @fusion_inner_long(%arg0: memref<10x2xf32>, %arg1: memref<10x10xf32>, %arg2: memref<10x2xf32>, %s: index) {
  %c0 = arith.constant 0 : index
  %cst_0 = arith.constant 1.000000e-03 : f32
  %c9 = arith.constant 9 : index
  %c10_i32 = arith.constant 10 : i32
  %c1_i32 = arith.constant 1 : i32
  %c100_i32 = arith.constant 100 : i32
  %c0_i32 = arith.constant 0 : i32
  %0 = memref.get_global @__constant_10x2xf32 : memref<10x2xf32>
  %1 = scf.for %arg3 = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg4 = %arg0) -> (memref<10x2xf32>)  : i32 {
    %alloc = memref.alloc() {alignment = 64 : i64} : memref<10xi32>
    affine.for %arg5 = 0 to 10 {
      %3 = arith.index_cast %arg5 : index to i32
      affine.store %3, %alloc[%arg5] : memref<10xi32>
    }
    %2 = scf.for %arg5 = %c0_i32 to %c10_i32 step %c1_i32 iter_args(%arg6 = %0) -> (memref<10x2xf32>)  : i32 {
      %alloc_5 = memref.alloc() : memref<2xf32>
      affine.for %arg7 = 0 to 2 {
        %16 = affine.load %arg4[%s, %arg7] : memref<10x2xf32>
        affine.store %16, %alloc_5[%arg7] : memref<2xf32>
      }
      %alloc_6 = memref.alloc() {alignment = 64 : i64} : memref<1x2xf32>
      affine.for %arg7 = 0 to 2 {
        %16 = affine.load %alloc_5[%arg7] : memref<2xf32>
        affine.store %16, %alloc_6[0, %arg7] : memref<1x2xf32>
      }
      %alloc_7 = memref.alloc() {alignment = 64 : i64} : memref<10x2xf32>
      affine.for %arg7 = 0 to 10 {
        affine.for %arg8 = 0 to 2 {
          %16 = affine.load %alloc_6[0, %arg8] : memref<1x2xf32>
          affine.store %16, %alloc_7[%arg7, %arg8] : memref<10x2xf32>
        }
      }
      %alloc_8 = memref.alloc() {alignment = 64 : i64} : memref<10x2xf32>
      affine.for %arg7 = 0 to 10 {
        affine.for %arg8 = 0 to 2 {
          %16 = affine.load %alloc_7[%arg7, %arg8] : memref<10x2xf32>
          %17 = affine.load %arg4[%arg7, %arg8] : memref<10x2xf32>
          %18 = arith.subf %16, %17 : f32
          affine.store %18, %alloc_8[%arg7, %arg8] : memref<10x2xf32>
        }
      }
      scf.yield %alloc_8 : memref<10x2xf32>
      // CHECK:      scf.for
      // CHECK:        scf.for
      // CHECK:          affine.for %{{.*}} = 0 to 10
      // CHECK-NEXT:       affine.for %{{.*}} = 0 to 2
      // CHECK-NOT:      affine.for
      // CHECK:          scf.yield
    }
    %alloc_2 = memref.alloc() {alignment = 64 : i64} : memref<10x2xf32>
    affine.for %arg5 = 0 to 10 {
      affine.for %arg6 = 0 to 2 {
        affine.store %cst_0, %alloc_2[%arg5, %arg6] : memref<10x2xf32>
      }
    }
    %alloc_3 = memref.alloc() {alignment = 64 : i64} : memref<10x2xf32>
    affine.for %arg5 = 0 to 10 {
      affine.for %arg6 = 0 to 2 {
        %3 = affine.load %alloc_2[%arg5, %arg6] : memref<10x2xf32>
        %4 = affine.load %2[%arg5, %arg6] : memref<10x2xf32>
        %5 = arith.mulf %3, %4 : f32
        affine.store %5, %alloc_3[%arg5, %arg6] : memref<10x2xf32>
      }
    }
    scf.yield %alloc_3 : memref<10x2xf32>
    // The nests above will be fused as well.
    // CHECK:      affine.for %{{.*}} = 0 to 10
    // CHECK-NEXT:   affine.for %{{.*}} = 0 to 2
    // CHECK-NOT:  affine.for
    // CHECK:      scf.yield
  }
  affine.for %arg3 = 0 to 10 {
    affine.for %arg4 = 0 to 2 {
      %2 = affine.load %1[%arg3, %arg4] : memref<10x2xf32>
      affine.store %2, %arg2[%arg3, %arg4] : memref<10x2xf32>
    }
  }
  return
}