<table border="1" cellspacing="0" cellpadding="8">
    <tr>
        <th>Issue</th>
        <td>
            <a href=https://github.com/llvm/llvm-project/issues/66697>66697</a>
        </td>
    </tr>

    <tr>
        <th>Summary</th>
        <td>
            [MLIR][Affine] --affine-parallelize failed to optimize some affine for loops
        </td>
    </tr>

    <tr>
      <th>Labels</th>
      <td>
            mlir
      </td>
    </tr>

    <tr>
      <th>Assignees</th>
      <td>
      </td>
    </tr>

    <tr>
      <th>Reporter</th>
      <td>
          Kuree
      </td>
    </tr>
</table>

<pre>
    When use `--affine-parallelize` with the following input, some inner `affine.for` operations do not get converted:
```mlir
module {
  func.func @bug(%arg1 : memref<1x16x16x1xi8>, %arg2 : memref<1x14x14x2xi32>) {
 %c-1_i32_2 = arith.constant -1 : i32
    %c127_i32_3 = arith.constant 127 : i32
    affine.for %arg3 = 0 to 1 {
      affine.for %arg4 = 0 to 14 {
        affine.for %arg5 = 0 to 14 {
          affine.for %arg6 = 0 to 2 {
            affine.for %arg7 = 0 to 3 {
              affine.for %arg8 = 0 to 3 {
                affine.for %arg9 = 0 to 1 {
 %0 = affine.load %arg1[%arg3, %arg4 + %arg7, %arg5 + %arg8, %arg9] : memref<1x16x16x1xi8>
                  %1 = affine.load %arg2[%arg3, %arg4, %arg5, %arg6] : memref<1x14x14x2xi32>
                  %2 = arith.extsi %0 : i8 to i32
                  %3 = arith.subi %2, %c-1_i32_2 : i32
                  %4 = arith.muli %3, %c127_i32_3 : i32
                  %5 = arith.addi %1, %4 : i32
 affine.store %5, %arg2[%arg3, %arg4, %arg5, %arg6] : memref<1x14x14x2xi32>
                }
              }
 }
          }
        }
      }
    }
    return
 }
}
```

Result:
```mlir
module {
  func.func @bug(%arg0: memref<1x16x16x1xi8>, %arg1: memref<1x14x14x2xi32>) {
 %c-1_i32 = arith.constant -1 : i32
    %c127_i32 = arith.constant 127 : i32
    affine.parallel (%arg2) = (0) to (1) {
      affine.parallel (%arg3) = (0) to (14) {
        affine.parallel (%arg4) = (0) to (14) {
          affine.parallel (%arg5) = (0) to (2) {
 affine.for %arg6 = 0 to 3 {
              affine.for %arg7 = 0 to 3 {
 affine.parallel (%arg8) = (0) to (1) {
                  %0 = affine.load %arg0[%arg2, %arg3 + %arg6, %arg4 + %arg7, %arg8] : memref<1x16x16x1xi8>
                  %1 = affine.load %arg1[%arg2, %arg3, %arg4, %arg5] : memref<1x14x14x2xi32>
                  %2 = arith.extsi %0 : i8 to i32
                  %3 = arith.subi %2, %c-1_i32 : i32
                  %4 = arith.muli %3, %c127_i32 : i32
 %5 = arith.addi %1, %4 : i32
                  affine.store %5, %arg1[%arg2, %arg3, %arg4, %arg5] : memref<1x14x14x2xi32>
 }
              }
            }
          }
        }
 }
    }
    return
  }
}
```

How to reproduce: https://godbolt.org/z/Gj61WPr89
</pre>
<img width="1px" height="1px" alt="" src="http://email.email.llvm.org/o/eJzMV12P2ygU_TX4BSUCbPzx4IdMs9ld7a606ksfK2zjhAqbCHBntr9-BU5qJyGeTNWROrI89s09h8P15SCYMWLfc14C-gToNmKDPShd_jVozqNKNf-Vnw68h4PhEKRotWJtK3q-OjLNpORSfOMgRfBZ2AO0Bw5bJaV6Fv0eiv44WEA-QKM6DkXfc-0YRvy6Vdrh1JFrZoXqDWwU7JWFe25hrfqvXFvegHgD0BagDUjReHVS6DHUqWaQHILsaXyHsB36eu1uECSoGvaA5IBQpvcYgngDO95p3oL4A37Bqb9eRA7i35zGMY_c5CXuIi8iJj6xmA0HCK1X-LOIyWeH20KmhT2sa9Uby3oLV-OoDnrSBz0Gk8yD4hAIk-wWNZXspHOEImgVxPMCBJOTWXJynR3Ip8v5AUQ6IUgIEIBkEyQOQwKg_AFQAFaEqwUIReMnGAFSseaEwG4tjJWemiOBgDyd5U9hOgvnU7gAdLvcdUHxvkXwHVkkKGsmZXpMQ8NfNvP98eftzF-sEedabaDIXRkvuvMGPm9sM1QeTU7S5itm8wpPMuPpBul5zhO_WEWvEdEZEWsaT4RPRMkV_lR1Y5XmHjtzh_cvP8i2oR-mcCDhJnQVuHi9eNHcDrq_Zp4ezp57evX3j9wM0v4kX0aP2TL-IVf-AU9-qyOfN0H4fUbEq4m3LoLcs1XuEV-KhIsc8R2OJECyQJO8iWaBiIaJyCXP0q7wuMXf2Rfuassfr_eVKdzzfvR9lZOpA-OZyaevbgn5O3g_Dsq640O_tPf_JOe_onmjzd_8Lfj-e5T-VaNfDi54_oN2_6jf_6Ge3VfX_KhVM9TcTe5g7dG4LYDsANntVVMpaddK7wHZfQNk9_uXFH_6V-dF1JRxU8QFi3iJ04KiDOM0jw5lkSFaZ0neMtw0eRJjltZFHrc0rhKa4TwSJUEkRgXOCUYpytd1RrMsy1JW46pgRQUSxDsm5FrKr50bOxLGDLxM07TIIskqLo0_0hDityZC3OFGly59VQ17AxIkhbFmIrDCSn8M-ufvPz8CugX0aeO7wn3U0MEHtkxI3rj6qKMVnQv5w86YCp2tSaWOJhq0LK-qJuxhqNa16gDZOQmnf6ujVl94bQHZ-QkZQHZ-Tv8HAAD__yWQgtM">