affine-loop-fusion does not respect fusion-compute-tolerance

[aqjune]

Hello all,

It seems -affine-loop-fusion does not respect its fusion-compute-tolerance option, sometimes causing a big increase in running time.

I attach the source file (a.mlir - a.zip) and its output (a_out.zip).
I ran:

bin/mlir-opt '-affine-loop-fusion=mode=producer fusion-compute-tolerance=0.0' a.mlir -o a_out.mlir

a.mlir contains three loops that look as follows:

affine.for %i = 0 to 131072 {
    ... // A
}
affine.for %j = 0 to 30 {
    affine.for %i = 0 to 131072 {
        ... // B
    }
}

affine-loop-fusion fuses the loops as follows:

affine.for %j = 0 to 30 {
    affine.for %i = 0 to 131072 {
        ... // A <- this iterates 29 times more than before!
        ... // B
    }
}

Any idea about how to fully disable this fusion?

Thanks in advance,

[llvmbot]

@llvm/issue-subscribers-mlir-affine

[aqjune]

If it is due to the existence of an unknown function call in // B, I wonder what is the proper way to properly control the behavior of affine loop fusion in this case.

[bondhugula]

With the debug flag -debug-only=affine-fusion, you can get a detailed log of why the fusion succeeded. Also, please don't zip the test case file - you can paste it inline since it's small enough. Adding the test case here.

#map = affine_map<(d0) -> (d0 mod 65536)>
module attributes {poly.he_param = "FVa"}  {
  memref.global "private" constant @prime_moduli_FVa : memref<30xi64> = dense<[2305843009146585089, 2251799756013569, 2251799787995137, 2251800352915457, 2251799780917249, 2251799666884609, 2251799678943233, 2251799696244737, 2251800082382849, 2251799776198657, 2251799929028609, 2251799774887937, 2251799849336833, 2251799883153409, 2251799777771521, 2251799879483393, 2251799772266497, 2251799763091457, 2251799844093953, 2251799823384577, 2251799851958273, 2251799789568001, 2251799797432321, 2251799799267329, 2251799836753921, 2251799806345217, 2251799807131649, 2251799818928129, 2251799816568833, 2251799815520257]>
  memref.global "private" constant @__constant_30xi64_1 : memref<30xi64> = dense<[536870904, 472446394368, 210453389312, 2247382979043329, 268435447808, 1202590834688, 1103806586880, 962072666112, 2249598911635457, 307090153472, 2250855036215297, 317827571712, 2251507791552513, 2251230799978497, 294205251584, 2251260861079553, 339302408192, 414464335872, 2251550735982593, 2251720366481409, 2251486319337473, 197568487424, 133143977984, 118111592448, 2251610858184705, 60129533952, 53687083008, 2251756869246977, 2251776194240513, 2251784783126529]>
  memref.global "private" constant @__constant_30xi64_0 : memref<30xi64> = dense<[4294967232, 3870280961785855, 1724034184912894, -36187118094876667, 2199023220443134, 9851624760016929, 9042384100794392, 7881299691831310, -18031988610955505, 2515682579122174, -7740561530159118, 2603643512324094, -2392537331269631, -4661929225058305, 2410129459415038, -4415638635231232, 2779565379035134, 3395291915749374, -2040693620703231, -650910947948544, -2568459185938431, 1618481066311678, 1090715475517439, 967570171529215, -1548112423157759, 492581143740415, 439804585281535, -351843787177984, -193514113349632, -123145369319424]>
  memref.global "private" constant @__constant_30xi64 : memref<30xi64> = dense<[8, 8192, 8192, 8191, 8192, 8192, 8192, 8192, 8191, 8192, 8191, 8192, 8191, 8191, 8192, 8191, 8192, 8192, 8191, 8191, 8191, 8192, 8192, 8192, 8191, 8192, 8192, 8191, 8191, 8191]>
  func.func @encode(%arg0: memref<65536xcomplex<f64>> {llvm.noalias}, %arg1: memref<30x131072xi64> {llvm.noalias}) {
    %c65536 = arith.constant 65536 : index
    %cst = arith.constant 0.000000e+00 : f64
    %cst_0 = arith.constant 0x4320000000380004 : f64
    %cst_1 = arith.constant 5.000000e-01 : f64
    %0 = memref.alloc() {alignment = 128 : i64} : memref<30x131072xi64>
    %1 = memref.alloc() {alignment = 128 : i64} : memref<131072xi1>
    %2 = memref.alloc() {alignment = 128 : i64} : memref<131072xi128>
    %3 = memref.get_global @__constant_30xi64 : memref<30xi64>
    %4 = memref.get_global @__constant_30xi64_0 : memref<30xi64>
    %5 = memref.get_global @__constant_30xi64_1 : memref<30xi64>
    %6 = memref.get_global @prime_moduli_FVa : memref<30xi64>
    func.call @__external_backward_fft(%arg0) : (memref<65536xcomplex<f64>>) -> ()
    affine.for %arg2 = 0 to 131072 {
      %7 = affine.apply #map(%arg2)
      %8 = affine.load %arg0[%7] : memref<65536xcomplex<f64>>
      %9 = arith.cmpi ult, %arg2, %c65536 : index
      %10 = complex.im %8 : complex<f64>
      %11 = complex.re %8 : complex<f64>
      %12 = arith.select %9, %11, %10 : f64
      %13 = arith.cmpf olt, %12, %cst : f64
      %14 = arith.negf %12 : f64
      %15 = arith.select %13, %14, %12 : f64
      %16 = arith.mulf %15, %cst_0 : f64
      %17 = arith.addf %16, %cst_1 : f64
      %18 = arith.fptosi %17 : f64 to i128
      affine.store %18, %2[%arg2] : memref<131072xi128>
      affine.store %13, %1[%arg2] : memref<131072xi1>
    }
    affine.for %arg2 = 0 to 30 {
      affine.for %arg3 = 0 to 131072 {
        %7 = affine.load %6[%arg2] : memref<30xi64>
        %8 = affine.load %3[%arg2] : memref<30xi64>
        %9 = affine.load %5[%arg2] : memref<30xi64>
        %10 = affine.load %4[%arg2] : memref<30xi64>
        %11 = affine.load %2[%arg3] : memref<131072xi128>
        %12 = affine.load %1[%arg3] : memref<131072xi1>
        %13 = func.call @__external_reduce_barrett(%7, %8, %9, %10, %11) {outputModFac = 1 : i64} : (i64, i64, i64, i64, i128) -> i64
        %14 = arith.subi %7, %13 : i64
        %15 = arith.select %12, %14, %13 : i64
        affine.store %15, %0[%arg2, %arg3] : memref<30x131072xi64>
      }
    }
    func.call @__external_levelwise_forward_ntt(%0) : (memref<30x131072xi64>) -> ()
    affine.for %arg2 = 0 to 30 {
      affine.for %arg3 = 0 to 131072 {
        %7 = affine.load %0[%arg2, %arg3] : memref<30x131072xi64>
        affine.store %7, %arg1[%arg2, %arg3] : memref<30x131072xi64>
      }
    }
    memref.dealloc %2 : memref<131072xi128>
    memref.dealloc %1 : memref<131072xi1>
    memref.dealloc %0 : memref<30x131072xi64>
    return
  }
  func.func private @__external_forward_ntt(memref<30x131072xi64>, i64)
  func.func private @__external_backward_ntt(memref<30x131072xi64>, i64)
  func.func private @__external_levelwise_forward_ntt(memref<30x131072xi64>)
  func.func private @__external_levelwise_backward_ntt(memref<30x131072xi64>)
  func.func private @__external_reduce_barrett(i64, i64, i64, i64, i128) -> i64
  func.func private @__external_sample_uniform(memref<30x131072xi64>)
  func.func private @__external_sample_gaussian_without_ntt(memref<30x131072xi64>)
  func.func private @__external_sample_ZO_without_ntt(memref<30x131072xi64>)
  func.func private @__external_forward_fft(memref<65536xcomplex<f64>>)
  func.func private @__external_backward_fft(memref<65536xcomplex<f64>>)
}

From the fusion log:

Evaluating src loop 12 for dst loop 15
Found valid slice for depth: 1
Found valid slice for depth: 2
Max legal depth for fusion: 2
Skipping profitability analysis. Not supported for this case.
Fused src loop 12 into dst loop 15 at depth 2:
affine.for %arg2 = 0 to 30 {
  affine.for %arg3 = 0 to 131072 {
    ...

There is an issue to fix here by extending to support multiple producer stores, without which we shouldn't have been going ahead with fusion.

[bondhugula]

Regardless of the support for multiple producer stores, we should have checking if the compute tolerance is going to be exceeded. This is easy to fix.

[bondhugula]

Fixed by #128454