[mlir] Inconsistent output when executing MLIR program with `--linalg-fuse-elementwise-ops` `linalg-fold-unit-extent-dims` `--affine-loop-fusion`

[Emilyaxe]

git version: b2b267e

system: Ubuntu 18.04.6 LTS

Description:

I am experiencing an inconsistent result when executing the same MLIR program with and without --linalg-fuse-elementwise-ops linalg-fold-unit-extent-dims --affine-loop-fusion .
The output also varies when running the process with optimizations multiple times.
The output becomes correct when any of these three options is removed, so I'm unsure which optimization contains the bug.

Steps to Reproduce:

1. MLIR Program (tosa.mlir):

tosa.mlir:

module {
  func.func private @printMemrefI32(tensor<*xi32>)
  func.func private @printMemrefF32(tensor<*xf32>)
  func.func @main() {
    %1 = "tosa.const"() <{values = dense<-7982> : tensor<1x22x10xi32>}> : () -> tensor<1x22x10xi32>
    %73 = tosa.argmax %1 {axis = 1 : i32} : (tensor<1x22x10xi32>) -> tensor<1x10xi32>
    %88 = tosa.abs %73 : (tensor<1x10xi32>) -> tensor<1x10xi32>
    %89 = tosa.clamp %88 {max_val = 16 : i32,  min_val = 0 : i32} : (tensor<1x10xi32>) -> tensor<1x10xi32>
    %90 = tosa.arithmetic_right_shift %73, %89 {round = true} : (tensor<1x10xi32>, tensor<1x10xi32>) -> tensor<1x10xi32>
    %cast = tensor.cast %90 : tensor<1x10xi32> to tensor<*xi32>
    call @printMemrefI32(%cast) : (tensor<*xi32>) -> ()
    return
  }
}

2. Command to Run without optimizations :

/data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt a.mlir \
-pass-pipeline="builtin.module(func.func(tosa-to-linalg-named,tosa-to-linalg))" | /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt  \
-tosa-to-arith --linalg-fuse-elementwise-ops --linalg-fold-unit-extent-dims="use-rank-reducing-slices"  \
  -one-shot-bufferize="bufferize-function-boundaries"  -finalize-memref-to-llvm   -convert-linalg-to-affine-loops \
 --expand-strided-metadata   -convert-vector-to-llvm  -convert-arith-to-llvm  -lower-affine \
 -convert-math-to-llvm   -finalize-memref-to-llvm  -convert-scf-to-cf  -convert-arith-to-llvm -convert-cf-to-llvm -convert-func-to-llvm  \
-reconcile-unrealized-casts | timeout 10 /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-cpu-runner  \
-e main -entry-point-result=void --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_c_runner_utils.so \
 --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_runner_utils.so \
--shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_async_runtime.so

3. Output without optimizations ::

[[0,   0,   0,   0,   0,   0,   0,   0,   0,   0]]

4. Command to Run with optimizations :

/data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt a.mlir \
-pass-pipeline="builtin.module(func.func(tosa-to-linalg-named,tosa-to-linalg))" | /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt  \
-tosa-to-arith --linalg-fuse-elementwise-ops --linalg-fold-unit-extent-dims="use-rank-reducing-slices"  \
  -one-shot-bufferize="bufferize-function-boundaries"  -finalize-memref-to-llvm   -convert-linalg-to-affine-loops \
--affine-loop-fusion  --expand-strided-metadata   -convert-vector-to-llvm  -convert-arith-to-llvm  -lower-affine \
 -convert-math-to-llvm   -finalize-memref-to-llvm  -convert-scf-to-cf  -convert-arith-to-llvm -convert-cf-to-llvm -convert-func-to-llvm  \
-reconcile-unrealized-casts | timeout 10 /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-cpu-runner  \
-e main -entry-point-result=void --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_c_runner_utils.so \
 --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_runner_utils.so \
--shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_async_runtime.so

5. Output with optimizations :

[[-17909,   -17909,   -17909,   -17909,   -17909,   -17909,   -17909,   -17909,   -17909,   -17909]]

[llvmbot]

@llvm/issue-subscribers-mlir-linalg

Author: None (Emilyaxe)

git version: 953838d

system: Ubuntu 18.04.6 LTS

Description:

I am experiencing an inconsistent result when executing the same MLIR program with and without --linalg-fuse-elementwise-ops linalg-fold-unit-extent-dims --affine-loop-fusion .
The output also varies when running the process with optimizations multiple times.
The output becomes correct when any of these three options is removed, so I'm unsure which optimization contains the bug.

Steps to Reproduce:

1. MLIR Program (tosa.mlir):

tosa.mlir:

module {
  func.func private @<!-- -->printMemrefI32(tensor&lt;*xi32&gt;)
  func.func private @<!-- -->printMemrefF32(tensor&lt;*xf32&gt;)
  func.func @<!-- -->main() {
    %1 = "tosa.const"() &lt;{value = dense&lt;-7982&gt; : tensor&lt;1x22x10xi32&gt;}&gt; : () -&gt; tensor&lt;1x22x10xi32&gt;
    %73 = tosa.argmax %1 {axis = 1 : i32} : (tensor&lt;1x22x10xi32&gt;) -&gt; tensor&lt;1x10xi32&gt;
    %88 = tosa.abs %73 : (tensor&lt;1x10xi32&gt;) -&gt; tensor&lt;1x10xi32&gt;
    %89 = tosa.clamp %88 {max_fp = 1.600000e+01 : f32, max_int = 16 : i64, min_fp = 0.000000e+00 : f32, min_int = 0 : i64} : (tensor&lt;1x10xi32&gt;) -&gt; tensor&lt;1x10xi32&gt;
    %90 = tosa.arithmetic_right_shift %73, %89 {round = true} : (tensor&lt;1x10xi32&gt;, tensor&lt;1x10xi32&gt;) -&gt; tensor&lt;1x10xi32&gt;
    %cast = tensor.cast %90 : tensor&lt;1x10xi32&gt; to tensor&lt;*xi32&gt;
    call @<!-- -->printMemrefI32(%cast) : (tensor&lt;*xi32&gt;) -&gt; ()
    return
  }
}

2. Command to Run without optimizations :

/data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt a.mlir -pass-pipeline="builtin.module(func.func(tosa-to-linalg-named,tosa-to-linalg))" | /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt  -tosa-to-arith  -one-shot-bufferize="bufferize-function-boundaries"  -finalize-memref-to-llvm   -convert-linalg-to-affine-loops  --expand-strided-metadata   -convert-vector-to-llvm  -convert-arith-to-llvm  -lower-affine  -convert-math-to-llvm   -finalize-memref-to-llvm  -convert-scf-to-cf   -convert-func-to-llvm  -reconcile-unrealized-casts | timeout 10 /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-cpu-runner -e main -entry-point-result=void --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_c_runner_utils.so --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_runner_utils.so --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_async_runtime.so

3. Output without optimizations ::

[[0,   0,   0,   0,   0,   0,   0,   0,   0,   0]]

4. Command to Run with optimizations :

/data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt a.mlir -pass-pipeline="builtin.module(func.func(tosa-to-linalg-named,tosa-to-linalg))" | /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt  -tosa-to-arith --linalg-fuse-elementwise-ops --linalg-fold-unit-extent-dims="use-rank-reducing-slices"    -one-shot-bufferize="bufferize-function-boundaries"  -finalize-memref-to-llvm   -convert-linalg-to-affine-loops  --affine-loop-fusion  --expand-strided-metadata   -convert-vector-to-llvm  -convert-arith-to-llvm  -lower-affine  -convert-math-to-llvm   -finalize-memref-to-llvm  -convert-scf-to-cf   -convert-func-to-llvm  -reconcile-unrealized-casts | timeout 10 /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-cpu-runner -e main -entry-point-result=void --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_c_runner_utils.so --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_runner_utils.so --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_async_runtime.so

5. Output with optimizations :

[[-17909,   -17909,   -17909,   -17909,   -17909,   -17909,   -17909,   -17909,   -17909,   -17909]]

[Emilyaxe]

I tried reproducing this issue using historical commit versions and found that the inconsistent results occur with commit 4bce27015 (which results in a random value), but not with the previous commit 6a3982f (which results in 0).
Hi @guy-david, I am sorry to disturb you, but could you please take a moment to review this problem? Thank you!

[guy-david]

Hi! Sorry for the late reply.
I couldn't reproduce the issue with the commit you've built from nor on main, unfortunately (I also had to slightly modify the input). Can you verify once more on the latest main branch and post the results?

The commit in question allows the hoisting of these LLVM operations to allow further optimizations to happen in parent operations. In the final optimized IR I see some extra llvm.zero and llvm.poison operations, so maybe you've uncovered an underlying UB that only now begins to show. I suggest to run the final IR and add prints for the intermediate results to locate the problematic computation.

$ git rev-parse HEAD
2cb5241c7724cbf4fa02b25547e098bfc635aba9

$ cat tosa.mlir 
module {
  func.func private @printMemrefI32(tensor<*xi32>)
  func.func private @printMemrefF32(tensor<*xf32>)
  func.func @main() {
    %1 = "tosa.const"() <{value = dense<-7982> : tensor<1x22x10xi32>}> : () -> tensor<1x22x10xi32>
    %73 = tosa.argmax %1 {axis = 1 : i32} : (tensor<1x22x10xi32>) -> tensor<1x10xi32>
    %88 = tosa.abs %73 : (tensor<1x10xi32>) -> tensor<1x10xi32>
    %89 = tosa.clamp %88 {max_fp = 1.600000e+01 : f32, max_int = 16 : i64, min_fp = 0.000000e+00 : f32, min_int = 0 : i64} : (tensor<1x10xi32>) -> tensor<1x10xi32>
    %90 = tosa.arithmetic_right_shift %73, %89 {round = true} : (tensor<1x10xi32>, tensor<1x10xi32>) -> tensor<1x10xi32>
    %cast = tensor.cast %90 : tensor<1x10xi32> to tensor<*xi32>
    call @printMemrefI32(%cast) : (tensor<*xi32>) -> ()
    return
  }
}

$ ./bin/mlir-opt tosa.mlir \
-pass-pipeline="builtin.module(func.func(tosa-to-linalg-named,tosa-to-linalg))" | ./bin/mlir-opt  \
-tosa-to-arith --linalg-fuse-elementwise-ops --linalg-fold-unit-extent-dims="use-rank-reducing-slices"  \
  -one-shot-bufferize="bufferize-function-boundaries"  -finalize-memref-to-llvm   -convert-linalg-to-affine-loops \
--affine-loop-fusion  --expand-strided-metadata   -convert-vector-to-llvm  -convert-arith-to-llvm  -lower-affine \
 -convert-math-to-llvm   -finalize-memref-to-llvm  -convert-scf-to-cf   -convert-func-to-llvm  \
-reconcile-unrealized-casts | timeout 10 /Users/guyda/projects/llvm-project/build/bin/mlir-runner  \
-e main -entry-point-result=void --shared-libs=/Users/guyda/projects/llvm-project/build/lib/libmlir_c_runner_utils.dylib \
 --shared-libs=/Users/guyda/projects/llvm-project/build/lib/libmlir_runner_utils.dylib \
--shared-libs=/Users/guyda/projects/llvm-project/build/lib/libmlir_async_runtime.dylib
Unranked Memref base@ = 0x600001c7c100 rank = 2 offset = 0 sizes = [1, 10] strides = [10, 1] data = 
[[0,   0,   0,   0,   0,   0,   0,   0,   0,   0]]

[Emilyaxe]

Hi @guy-david,

Sorry for the delayed response. I have thoroughly analyzed the issue and found that the root cause may lie with --affine-loop-fusion.
Following your suggestion, I used the version 2cb5241 to reproduce the inconsistency. Since the mlir-opt of 2cb5241 no longer supports the syntax for tosa.clamp, I made a slight modification to the input as follows:

module {
  func.func private @printMemrefI32(tensor<*xi32>)
  func.func private @printMemrefF32(tensor<*xf32>)
  func.func @main() {
    %1 = "tosa.const"() <{value = dense<0> : tensor<1x22x10xi32>}> : () -> tensor<1x22x10xi32>
    %73 = tosa.argmax %1 {axis = 1 : i32} : (tensor<1x22x10xi32>) -> tensor<1x10xi32>
    %88 = tosa.abs %73 : (tensor<1x10xi32>) -> tensor<1x10xi32>
    %cast = tensor.cast %88 : tensor<1x10xi32> to tensor<*xi32>
    call @printMemrefI32(%cast) : (tensor<*xi32>) -> ()
    return
  }
}

Running the input with the following command will trigger a random result.

/data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt-2cb5241c77 a.mlir \
-pass-pipeline="builtin.module(func.func(tosa-to-linalg-named,tosa-to-linalg))" | /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-opt-2cb5241c77  \
-tosa-to-arith --linalg-fuse-elementwise-ops --linalg-fold-unit-extent-dims="use-rank-reducing-slices"  \
  -one-shot-bufferize="bufferize-function-boundaries"  -finalize-memref-to-llvm   -convert-linalg-to-affine-loops \
--affine-loop-fusion  --expand-strided-metadata   -convert-vector-to-llvm    -lower-affine \
 -convert-math-to-llvm   -finalize-memref-to-llvm  -convert-scf-to-cf  -convert-cf-to-llvm -convert-arith-to-llvm  -convert-func-to-llvm  \
-reconcile-unrealized-casts | timeout 10 /data/szy/MLIR/llvm-release/llvm-project/build/bin/mlir-cpu-runner  \
-e main -entry-point-result=void --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_c_runner_utils.so \
 --shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_runner_utils.so \
--shared-libs=/data/szy/MLIR/llvm-release/llvm-project/build/lib/libmlir_async_runtime.so

Unranked Memref base@ = 0x5578bb8bf300 rank = 2 offset = 0 sizes = [1, 10] strides = [10, 1] data = 
[[1149217184,   1149217184,   1149217184,   1149217184,   1149217184,   1149217184,   1149217184,   1149217184,   1149217184,   1149217184]]

I printed the IR after the transformation of each pass and found that the IR becomes incorrect after running with --affine-loop-fusion.
The input IR (before running the --affine-loop-fusion pass) input.txt.
The output IR (after running the --affine-loop-fusion pass ) output.txt.

I debugged the --affine-loop-fusion pass (using the debug-only option) and discovered that during the optimization, it actually prints the correct IR. However, for some reason, the final IR becomes incorrect after the complete --affine-loop-fusion pass is executed.

HI @bondhugula, I notice that you have worked on this pass before, but could you please kindly take a moment to review this problem? Thank you!