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

    <tr>
        <th>Summary</th>
        <td>
            [MLIR][affine] Illegal affine loop fusion with vector types
        </td>
    </tr>

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

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

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

<pre>
    Hi, I was playing around with affine fusion with affine.vector_load/store operations inside loop and came across this particular case:

`func.func @main(%a: memref<64x512xf32>, %b: memref<64x512xf32>, %c: memref<64x512xf32>, %d: memref<64x4096xf32>, %e: memref<64x4096xf32>) {

    affine.for %j = 0 to 8 {
        %lhs = affine.vector_load %a[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
        %rhs = affine.vector_load %b[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
 %res = arith.addf %lhs, %rhs : vector<64x64xf32>
 affine.vector_store %res, %c[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
    }

    affine.for %j = 0 to 8 {
 %lhs = affine.vector_load %c[0, 0] : memref<64x512xf32>, vector<64x512xf32>
        %rhs = affine.vector_load %d[0, %j * 512] : memref<64x4096xf32>, vector<64x512xf32>
        %res = arith.subf %lhs, %rhs : vector<64x512xf32>
        affine.vector_store %res, %d[0, %j * 512] : memref<64x4096xf32>, vector<64x512xf32>
    }

 func.return
}`

Upon invoking affine-fusion on this IR with the following command:

`mlir-opt --pass-pipeline='builtin.module(affine-loop-fusion)' test.mlir`


I see that the loops are getting fused as follows:

` func.func @main(%arg0: memref<64x512xf32>, %arg1: memref<64x512xf32>, %arg2: memref<64x512xf32>, %arg3: memref<64x4096xf32>, %arg4: memref<64x4096xf32>) {
    %c0 = arith.constant 0 : index
    %alloc = memref.alloc() : memref<1x1xf32>
    %c0_0 = arith.constant 0 : index
    affine.for %arg5 = 0 to 8 {
      %0 = affine.vector_load %arg0[0, %c0 * 64] : memref<64x512xf32>, vector<64x64xf32>
      %1 = affine.vector_load %arg1[0, %c0 * 64] : memref<64x512xf32>, vector<64x64xf32>
      %2 = arith.addf %0, %1 : vector<64x64xf32>
      affine.vector_store %2, %arg2[0, %c0 * 64] : memref<64x512xf32>, vector<64x64xf32>
      %3 = affine.vector_load %arg0[0, %c0_0 * 64] : memref<64x512xf32>, vector<64x64xf32>
      %4 = affine.vector_load %arg1[0, %c0_0 * 64] : memref<64x512xf32>, vector<64x64xf32>
      %5 = arith.addf %3, %4 : vector<64x64xf32>
      affine.vector_store %5, %alloc[0, 0] : memref<1x1xf32>, vector<64x64xf32>
      %6 = affine.vector_load %alloc[0, 0] : memref<1x1xf32>, vector<64x512xf32>
      %7 = affine.vector_load %arg3[0, %arg5 * 512] : memref<64x4096xf32>, vector<64x512xf32>
      %8 = arith.subf %6, %7 : vector<64x512xf32>
      affine.vector_store %8, %arg3[0, %arg5 * 512] : memref<64x4096xf32>, vector<64x512xf32>
    }
 return
  }`
  
  
  Fusion here is illegal since  2nd loop can only be executed once 1st loop completely finishes all its iterations and produces the result to be consumed in 2nd loop. I'm attaching the llvm discourse thread here for further reading about the issue.
  
 https://discourse.llvm.org/t/affine-fusion-legality-for-vector-types/83079
 
  
 As per the discussion,  we should have a bailout in the presence of different-sized element types in the producer/consumer validity checking.
 
</pre>
<img width="1px" height="1px" alt="" src="http://email.email.llvm.org/o/eJy0WFFv6jgT_TXmZQRynITAAw_t5UNfpd2XK-3zlRNPiLuOHdlOW_bXr-wECrQFerflIUjOxGdmznjOJNw5udWIK5Lfk3w94b1vjF0Jy7dGC8srbneO0mJSGrFb_V8S9gMe4Jk76BTfSb0Fbk2vBTxL3wCva6kR6t5Jo4-XZk9YeWN_KcMFYRvnjUUwHVrupdEOpHZSIChjOuBaQMVbBF5Z4xz4RjrouPWy6hW3UHGHJL0jdE3o_jqnda-rWbgAyWjLpSZsQVjOSXoHLbYWa5L-mGcvecJe6pSR9H8hFMLy8qpFddVCnFtkdDk_NcGLJksgxf1xRACwT11tbNjgEUi6BgrewOLVGsYfYblqXDR5m3GIicjv6ejLIxB2B_OM5Gu4HNuwyXBnnu1vnCPbi8jllyMHSBwhrfTNjAtRjykYgQaX7i5uc-ruUJPD1gfivyFnpFj_FtFXGT54Sz_n49GdzxEr3qQnT9h72Gen4VbwE4pdX95A8UfbXaX6u2I5Zzu2KYu-t3pcL9ZkTo9t_uqMBqmfzN-xvUbPp2NLNXpohw8_h-7qG4TaKGWeg21l2pZr8bY5tkraqek8TKcdd27ayQ6V1EjSNWFF2UvlpZ61RvQKCVuMmKEbj8CELQkrwKPzs7DZmcvD9QEcIviG--hXeNwBtwhb9D74V_cOBXA3uuzeOgoftXG7pVe7MLfb5BYjdotRer2lc7vNbu3qY01X9KikK6Od59oDjVUmtcCXE2uulKniAwPELC7EjCxPCzN5Sd6WXoD7dTvgaRfidptfUBzCcnpJbQJfrycqhP1VikNYnlxGTr4Pmb0jOnuo5KrgXGpF7Lg-v83_9DOc_fpa7OwzrH0xdv4Ob-kIlv0n3vI9b_FofqjARwf0Rpfnl9L1u2Dv6yNheXGZnPSInKEzfKHaE5Yv3lH6-YhX3KjyHxC0OO7p3xbFQefhWNyH9b1WApz_bwZVb9AiSAdSKdxyBU7qCgGYFsMbUcWD8qsdlAj4glXvUYAJNonzo4lpO4Ue1Q5qqaVr0AFXCqR3IP3hPSu8W3XWiL5CF0XaouuVDy2-RAj60LcoQOoD-AweCCta4N7zqgkqHrVdPbUgpKtMb13QfItcDHEE_ah76xu0EFbjEFOafhgKpHM9zs7y0HjfxWGAbQjbHLadBZSZsVvCNp6wzckoNI2pkn43rY2dDvRM_a4LE91mkdJiuZ-bT7HuHHRooy8BqHfDePMD4BnBNaZXAhr-hMCh5FIFv6WO5p1FhyHppgYh6xotaj918h8UgApb1B6iA68PxERbwjZjZi08cSWF9DuoGqzCgLfPxUSsUrFMl3yCq6RIk4wu0iSZNKs0X9L5cimqgopFmS8rWuZI63SOtEiSjE7kilGWJUnCEprlSTIr2FxwUWR1scwSpEgyii2X6pDPSWRhlST5IltOFC9RufgBgLE43LFwIiZ2FeynZb91JKNKOu9ed_DSq_jR4M8_Hn6SfE3y-4GecJYexjoePwfEAj3-JjCwNSRr0lu1Oi2ArfRNX84q0xK2CYjj37Sz5hGrUAoxgMD0GMPTiv0bAAD__9340GM">