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

    <tr>
        <th>Summary</th>
        <td>
            [mlir] [affine] affine-loop-fusion incorrectly fuses loops
        </td>
    </tr>

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

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

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

<pre>
    Consider the following MLIR code snippet:

```mlir
func.func @foo(%input: memref<1x40x32x128xf32>, 
                                  %output: memref<1x32x40x128xf32>,
                                  %scale: memref<1x1x40x128xf32>) {
  affine.for %i0 = 0 to 1 {
    affine.for %i1 = 0 to 32 {
      affine.for %i2 = 0 to 40 {
        affine.for %i3 = 0 to 128 {
          %val = affine.load %input[%i0, %i2, %i1, %i3] : memref<1x40x32x128xf32>
          affine.store %val, %output[%i0, %i1, %i2, %i3] : memref<1x32x40x128xf32>
        }
      }
    }
  }

  affine.for %i0 = 0 to 1 {
    affine.for %i1 = 0 to 32 {
      affine.for %i2 = 0 to 40 {
        affine.for %i3 = 0 to 128 {
          %val = affine.load %output[%i0, %i1, %i2, %i3] : memref<1x32x40x128xf32>
          %scale_val = affine.load %scale[%i0, 0, %i2, %i3] : memref<1x1x40x128xf32>
          %result = arith.mulf %val, %scale_val : f32
          affine.store %result, %output[%i0, %i1, %i2, %i3] : memref<1x32x40x128xf32>
        }
      }
    }
  }
  return
}
```

the result produced by `mlir-opt --affine-loop-fusion` is:

```mlir
module {
  func.func @foo(%arg0: memref<1x40x32x128xf32>, %arg1: memref<1x32x40x128xf32>, %arg2: memref<1x1x40x128xf32>) {
    %c0 = arith.constant 0 : index // ❌ should not have this const 
    affine.for %arg3 = 0 to 1 {
      affine.for %arg4 = 0 to 32 {
        affine.for %arg5 = 0 to 40 {
          affine.for %arg6 = 0 to 128 {
            %0 = affine.load %arg0[%c0, %arg5, %arg4, %arg6] : memref<1x40x32x128xf32> // ❌ Uses %c0 instead of %arg3
            affine.store %0, %arg1[%c0, %arg4, %arg5, %arg6] : memref<1x32x40x128xf32> // ❌ Uses %c0 instead of %arg3
            %1 = affine.load %arg1[%arg3, %arg4, %arg5, %arg6] : memref<1x32x40x128xf32>
            %2 = affine.load %arg2[%arg3, 0, %arg5, %arg6] : memref<1x1x40x128xf32>
            %3 = arith.mulf %1, %2 : f32
            affine.store %3, %arg1[%arg3, %arg4, %arg5, %arg6] : memref<1x32x40x128xf32>
          }
        }
      }
    }
    return
  }
}
```

but the correct loop should be:

```mlir
affine.for %arg3 = 0 to 1 {
  ...
  %0 = affine.load %arg0[%arg3, %arg5, %arg4, %arg6]  // Should use %arg3
  affine.store %0, %arg1[%arg3, %arg4, %arg5, %arg6]  // Should use %arg3
  %1 = affine.load %arg1[%arg3, %arg4, %arg5, %arg6]  
  ...
}
```

Here is a godbolt link for reproduction [https://godbolt.org/z/K6nzjd4d5](https://godbolt.org/z/K6nzjd4d5). 
</pre>
<img width="1" height="1" alt="" src="http://email.email.llvm.org/o/eJzcV1uvqzYT_TXOyyjIjHEuDzzs7Jzo-9T2pVWfKy4m8amDkW12s8-vrzAkIUAuUs9RpW5FO0ZZnjWembWAxFq5L4WICd8Qvp0ltTtoE6d1mipRyUrMUp1_xu-6tDIXBtxBQKGV0n_Jcg-__Pz_XyHTuQBbyqoSjrA3Qv1nQdvPUUlD6FtRl1nQ_AMS0UJrgiuCXJZV3eyBozgaURD2Hp4iemJ4CnF1KhgS9oXgOxD6Bk__CHJdu3FAhk3Mm4CvxrNZosQgXDiMtgay3PiISVHIUgSFNs1mSYGwLVBwGsILZoQKryiGPdgIiFdgRG-AIyjrMeNqgPUn-0iUx3QblU5yuDSEb3z6vvIN8XkRnheM8C0861ufsaOxThvR0XexupYNKMMR9xTlqLNXSrLcXq6u6_Oq_f6P9OwHVfAqgD_uELfq6PGOJ2aKcyigAaURtlau5TPSHYJjrYrbmeln9QZNmEez1gb898cNwAhXm7KZvHb-zh7ZzmJjrd3hK6PzOhM5pJ_QeehcVw7m8_Zwc6V1NS9qK3VJFhSkvWu8R53XSlzGadqHE7Onr9iwR4bP_bVD4uvW6Vuf0V7bM11al5QOqO-DLHNxAoI7gjsgX5Cst2T1Dvaga5VDqR0ckg8B7iAt-K0wKd3E7Nm0xCeQ0V2ZT4D5A6lPwBeP5e7rQadU53vl5zej11Lz6zK6Lhcv-PS4pr9bYbt2yNI6keSgi3PxBlkOxUb7czLKMppMeCrL4VD9sywJ8vBOLbss_abvkOeYGO8Q4y3xdC-nKB_5p2dkE-Z5djictMxxH9mojz-kQn3ffOaiPf-83MGnfDStnX9MzbQxInPQeOXZJ1Jx1yhf84kgCFr-p-q8rdcDfZ5H-7c2xdqKmyl-rrCXO_OE6TuqBHrFmmzS_4QRIC0ksNd5qpUDJcs_oam-Ee3tz0ldAuGbg3OVv7_55Dt4oM2e4O4bwd1Pi_Lb1zzKOeFbgqvX4bgOYJbHLF-zdTITcbjkEVuskbHZIS44x7TgixVyFLxYcZGyJMnX-XqdsDyNZjJGipyuKKdRtAxZUCzZYsUEC5dFFC2SjERUHBOpAqU-jk0CM2ltLeKQI10uZypJhbL-vQvRzyBi8wZm4gY_T-u9JRFV0jp7jeCkU_5dzW9oxMY3bbeai_HDAciyk4H6hKJuLLP51c5qo-JBpaQ71GmQ6SPBXUPYfc0ro7-KzBHc-fwtwV13hI8Y_w4AAP__3-fOqg">