[mlir] [affine] affine-loop-fusion incorrectly fuses loops

[bubblepipe]

Consider the following MLIR code snippet:

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:

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:

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.

[llvmbot]

@llvm/issue-subscribers-mlir

Author: bubblepipe (bubblepipe)

Consider the following MLIR code snippet:

func.func @<!-- -->foo(%input: memref&lt;1x40x32x128xf32&gt;, 
                                  %output: memref&lt;1x32x40x128xf32&gt;,
                                  %scale: memref&lt;1x1x40x128xf32&gt;) {
  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&lt;1x40x32x128xf32&gt;
          affine.store %val, %output[%i0, %i1, %i2, %i3] : memref&lt;1x32x40x128xf32&gt;
        }
      }
    }
  }

  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&lt;1x32x40x128xf32&gt;
          %scale_val = affine.load %scale[%i0, 0, %i2, %i3] : memref&lt;1x1x40x128xf32&gt;
          %result = arith.mulf %val, %scale_val : f32
          affine.store %result, %output[%i0, %i1, %i2, %i3] : memref&lt;1x32x40x128xf32&gt;
        }
      }
    }
  }
  return
}

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

module {
  func.func @<!-- -->foo(%arg0: memref&lt;1x40x32x128xf32&gt;, %arg1: memref&lt;1x32x40x128xf32&gt;, %arg2: memref&lt;1x1x40x128xf32&gt;) {
    %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&lt;1x40x32x128xf32&gt; // ❌ Uses %c0 instead of %arg3
            affine.store %0, %arg1[%c0, %arg4, %arg5, %arg6] : memref&lt;1x32x40x128xf32&gt; // ❌ Uses %c0 instead of %arg3
            %1 = affine.load %arg1[%arg3, %arg4, %arg5, %arg6] : memref&lt;1x32x40x128xf32&gt;
            %2 = affine.load %arg2[%arg3, 0, %arg5, %arg6] : memref&lt;1x1x40x128xf32&gt;
            %3 = arith.mulf %1, %2 : f32
            affine.store %3, %arg1[%arg3, %arg4, %arg5, %arg6] : memref&lt;1x32x40x128xf32&gt;
          }
        }
      }
    }
    return
  }
}

but the correct loop should be:

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.

[CoTinker]

This is a producer-consumer fusion, and the outcome matches expectations.
You can see example in affine-loop-fusion:

llvm-project/mlir/include/mlir/Dialect/Affine/Passes.td

Lines 65 to 107 in 95d9e8a

	Example 1: Producer-consumer fusion.
	Input:
	```mlir
	func.func @producer_consumer_fusion(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
	%0 = memref.alloc() : memref<10xf32>
	%1 = memref.alloc() : memref<10xf32>
	%cst = arith.constant 0.000000e+00 : f32
	affine.for %arg2 = 0 to 10 {
	affine.store %cst, %0[%arg2] : memref<10xf32>
	affine.store %cst, %1[%arg2] : memref<10xf32>
	}
	affine.for %arg2 = 0 to 10 {
	%2 = affine.load %0[%arg2] : memref<10xf32>
	%3 = arith.addf %2, %2 : f32
	affine.store %3, %arg0[%arg2] : memref<10xf32>
	}
	affine.for %arg2 = 0 to 10 {
	%2 = affine.load %1[%arg2] : memref<10xf32>
	%3 = arith.mulf %2, %2 : f32
	affine.store %3, %arg1[%arg2] : memref<10xf32>
	}
	return
	}
	```
	Output:
	```mlir
	func.func @producer_consumer_fusion(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
	%0 = memref.alloc() : memref<1xf32>
	%1 = memref.alloc() : memref<1xf32>
	%cst = arith.constant 0.000000e+00 : f32
	affine.for %arg2 = 0 to 10 {
	affine.store %cst, %0[0] : memref<1xf32>
	affine.store %cst, %1[0] : memref<1xf32>
	%2 = affine.load %1[0] : memref<1xf32>
	%3 = arith.mulf %2, %2 : f32
	affine.store %3, %arg1[%arg2] : memref<10xf32>
	%4 = affine.load %0[0] : memref<1xf32>
	%5 = arith.addf %4, %4 : f32
	affine.store %5, %arg0[%arg2] : memref<10xf32>
	}
	return
	}
	```

[llvmbot]

@llvm/issue-subscribers-mlir-affine

Author: bubblepipe (bubblepipe)

Consider the following MLIR code snippet:

func.func @<!-- -->foo(%input: memref&lt;1x40x32x128xf32&gt;, 
                                  %output: memref&lt;1x32x40x128xf32&gt;,
                                  %scale: memref&lt;1x1x40x128xf32&gt;) {
  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&lt;1x40x32x128xf32&gt;
          affine.store %val, %output[%i0, %i1, %i2, %i3] : memref&lt;1x32x40x128xf32&gt;
        }
      }
    }
  }

  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&lt;1x32x40x128xf32&gt;
          %scale_val = affine.load %scale[%i0, 0, %i2, %i3] : memref&lt;1x1x40x128xf32&gt;
          %result = arith.mulf %val, %scale_val : f32
          affine.store %result, %output[%i0, %i1, %i2, %i3] : memref&lt;1x32x40x128xf32&gt;
        }
      }
    }
  }
  return
}

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

module {
  func.func @<!-- -->foo(%arg0: memref&lt;1x40x32x128xf32&gt;, %arg1: memref&lt;1x32x40x128xf32&gt;, %arg2: memref&lt;1x1x40x128xf32&gt;) {
    %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&lt;1x40x32x128xf32&gt; // ❌ Uses %c0 instead of %arg3
            affine.store %0, %arg1[%c0, %arg4, %arg5, %arg6] : memref&lt;1x32x40x128xf32&gt; // ❌ Uses %c0 instead of %arg3
            %1 = affine.load %arg1[%arg3, %arg4, %arg5, %arg6] : memref&lt;1x32x40x128xf32&gt;
            %2 = affine.load %arg2[%arg3, 0, %arg5, %arg6] : memref&lt;1x1x40x128xf32&gt;
            %3 = arith.mulf %1, %2 : f32
            affine.store %3, %arg1[%arg3, %arg4, %arg5, %arg6] : memref&lt;1x32x40x128xf32&gt;
          }
        }
      }
    }
    return
  }
}

but the correct loop should be:

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.