[CPU] Incorrect results depending on `native_vector_size` for fusion of 2 generics.

[bjacob]

See attached reproduction. Unzip, cd into that directory and run ./repro.sh

repro.zip

repro.sh is just compiling, running, and print the number of NaNs. The expected result is 0 (even though NaNs occur as intermediate results, they should not occur in the final outputs). The actual value is 311. If you edit the testcase to change native_vector_size from 32 to 64, then the printed value changes from 311 to 0, which confirms that there is a bug: the outputs should not depend on the value of native_vector_size.

This testcase is already minimal in the following sense:

1. The source has to be post Flow. If we just make the source be pure linalg, it doesn't reproduce anymore, presumably because it is essential to have one dispatch fusing the two generics here.
2. If we try pre-evaluating the first generic (reduction) and feeding its output to just the second generic, this stops reproducing too. Again, this seems to show that the problem is specifically about the fusion.
3. The testcase also stops reproducing if I try removing any of the remaining ops in the second generic.

[bjacob]

This is a bug with vector masking. The simplest proof / temporary unblocking fix is to disable vector masking, e.g.:

--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
@@ -142,9 +142,7 @@ void LLVMCPULowerExecutableTargetPass::runOnOperation() {
   pipelineOpts.decomposePackUnPackOps =
       isOptEnabled(funcOp, getEnableDecompositionStr());
   pipelineOpts.lowerToAVX2 = hasAVX2Feature(targetConfig);
-  pipelineOpts.enableVectorMasking =
-      isX86(targetConfig) || isRISCV(targetConfig) ||
-      (isAArch64(targetConfig) && hasAnySVEFeature(targetConfig));
+  pipelineOpts.enableVectorMasking = false;
   // TODO(#16956): The decomposition of attention op leads to complex control
   // flow, which leads to non-trivial stack allocation. Enforcing masking for
   // targets that do not support native vector masking enables the

The step in mlir-print-ir-after-all where the bug happens is:

// -----// IR Dump After GenericVectorizationPass (iree-codegen-generic-vectorization) //----- //
[...]
      %13 = scf.for %arg4 = %c0 to %c625 step %c8 iter_args(%arg5 = %12) -> (tensor<8xf32>) {
        %15 = affine.min affine_map<(d0) -> (-d0 + 625, 8)>(%arg4)
        %16 = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%arg2, %arg0]
        %extracted_slice_3 = tensor.extract_slice %5[%16, %arg4] [8, %15] [1, 1] : tensor<1600x625xf32> to tensor<8x?xf32>
        %17 = vector.create_mask %c8, %15 : vector<8x8xi1>
        %18 = vector.transfer_read %extracted_slice_3[%c0, %c0], %0, %17 {in_bounds = [true, true]} : tensor<8x?xf32>, vector<8x8xf32>
        %19 = vector.transfer_read %arg5[%c0], %0 {in_bounds = [true]} : tensor<8xf32>, vector<8xf32>
        %20 = vector.mask %17 { vector.multi_reduction <add>, %18, %19 [1] : vector<8x8xf32> to vector<8xf32> } : vector<8x8xi1> -> vector<8xf32>
        %21 = vector.transfer_write %20, %arg5[%c0] {in_bounds = [true]} : vector<8xf32>, tensor<8xf32>
        scf.yield %21 : tensor<8xf32>
      }

As 625 is not a multiple of 8, the last iteration has %extracted_slice_3 's inner dimension smaller than 8, which is OK, it's why it's a dynamic dimension. But then the %18 = vector.transfer_read %extracted_slice_3 incorrectly has in_bounds = [true, true], which should be in_bounds = [true, false]. This is probably coming from an upstream vectorization pattern. The proper fix is probably to fix that, upstream.

Regarding why native_vector_size = 64 seems to "fix" it, well it's still making an out-of-bounds access, it's just being lucky:

// -----// IR Dump After GenericVectorizationPass (iree-codegen-generic-vectorization) //----- //
[...]
      %13 = scf.for %arg4 = %c0 to %c625 step %c16 iter_args(%arg5 = %12) -> (tensor<4xf32>) {
        %15 = affine.min affine_map<(d0) -> (-d0 + 625, 16)>(%arg4)
        %16 = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%arg2, %arg0]
        %extracted_slice_3 = tensor.extract_slice %5[%16, %arg4] [4, %15] [1, 1] : tensor<1600x625xf32> to tensor<4x?xf32>
        %17 = vector.create_mask %c4, %15 : vector<4x16xi1>
        %18 = vector.transfer_read %extracted_slice_3[%c0, %c0], %0, %17 {in_bounds = [true, true]} : tensor<4x?xf32>, vector<4x16xf32>
        %19 = vector.transfer_read %arg5[%c0], %0 {in_bounds = [true]} : tensor<4xf32>, vector<4xf32>
        %20 = vector.mask %17 { vector.multi_reduction <add>, %18, %19 [1] : vector<4x16xf32> to vector<4xf32> } : vector<4x16xi1> -> vector<4xf32>
        %21 = vector.transfer_write %20, %arg5[%c0] {in_bounds = [true]} : vector<4xf32>, tensor<4xf32>
        scf.yield %21 : tensor<4xf32>
      }

Exact same issue, kind of worse actually because instead of reading 7 out-of-bounds elements it's now reading 15. Luck.

Tried --iree-llvmcpu-sanitize=address as that would have been a good use case for it, but sadly it didn't see the out-of-bounds access; perhaps because the undefined behavior morphs into something else that still corrupts results in a later lowering.

[bjacob]

Here is a minimized MLIR testcase that can be vectorized either using iree-codegen-generic-vectorization (reproducing the bug) or with transform dialect in mlir-opt but there, it looks different.

The difference is that while both do vector.create_mask, iree-codegen-generic-vectorization`` only uses vector.maskon thevector.multi_reductionbut leaves thetransfer_readunmasked, whilemlir-optappliesvector.maskto thetransfer_read` op.

I'm not familiar with vector.mask semantics at all, but it sounds like it somehow makes in_bounds = true to be okay, so IIUC, the transform-dialect flavor might be okay but iree-codegen-generic-vectorization is wrong ?

The testcase:

func.func @bad_masking_repro(%m : tensor<1600x625xf32>) -> tensor<1600x25x25xf32> {
  %c8 = arith.constant 8 : index
  %c625 = arith.constant 625 : index
  %cst = arith.constant 0.000000e+00 : f32
  %c0 = arith.constant 0 : index
  %empty = tensor.empty() : tensor<8xf32>
  %filled = linalg.fill ins(%cst : f32) outs(%empty : tensor<8xf32>) -> tensor<8xf32>
  %reduced = scf.for %iv = %c0 to %c625 step %c8 iter_args(%acc = %filled) -> (tensor<8xf32>) {
    %w = affine.min affine_map<(d0) -> (-d0 + 625, 8)>(%iv)
    %tile = tensor.extract_slice %m[%c0, %iv] [8, %w] [1, 1] : tensor<1600x625xf32> to tensor<8x?xf32>
    %next = linalg.generic {
      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],
      iterator_types = ["parallel", "reduction"]
    } ins(%tile : tensor<8x?xf32>) outs(%acc : tensor<8xf32>) {
    ^bb0(%in: f32, %out: f32):
      %s = arith.addf %out, %in : f32
      linalg.yield %s : f32
    } -> tensor<8xf32>
    scf.yield %next : tensor<8xf32>
  }
  %e = tensor.extract %reduced[%c0] : tensor<8xf32>
  %base = tensor.empty() : tensor<1600x25x25xf32>
  %out = tensor.insert %e into %base[%c0, %c0, %c0] : tensor<1600x25x25xf32>
  return %out : tensor<1600x25x25xf32>
}

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
    %g = transform.structured.match ops{["linalg.generic"]} in %root : (!transform.any_op) -> !transform.any_op
    transform.structured.vectorize %g vector_sizes [8, 8] : !transform.any_op
    transform.yield
  }
}

Reproduce the bug with iree-codegen-generic-vectorization :

iree-opt a.mlir --pass-pipeline="builtin.module(func.func(iree-codegen-generic-vectorization{enable-vector-masking=true fold-cast-into-contract=true}))"

Result:

    %3 = scf.for %arg1 = %c0 to %c625 step %c8 iter_args(%arg2 = %2) -> (tensor<8xf32>) {
      %5 = affine.min #map(%arg1)
      %extracted_slice = tensor.extract_slice %arg0[%c0, %arg1] [8, %5] [1, 1] : tensor<1600x625xf32> to tensor<8x?xf32>
      %6 = vector.create_mask %c8, %5 : vector<8x8xi1>
      %7 = vector.transfer_read %extracted_slice[%c0, %c0], %0, %6 {in_bounds = [true, true]} : tensor<8x?xf32>, vector<8x8xf32>
      %8 = vector.transfer_read %arg2[%c0], %0 {in_bounds = [true]} : tensor<8xf32>, vector<8xf32>
      %9 = vector.mask %6 { vector.multi_reduction <add>, %7, %8 [1] : vector<8x8xf32> to vector<8xf32> } : vector<8x8xi1> -> vector<8xf32>
      %10 = vector.transfer_write %9, %arg2[%c0] {in_bounds = [true]} : vector<8xf32>, tensor<8xf32>
      scf.yield %10 : tensor<8xf32>
    }

Vectorize using Transform dialect in mlir-opt:

mlir-opt a.mlir -transform-interpreter -canonicalize -cse

Result:

    %3 = scf.for %arg1 = %c0 to %c625 step %c8 iter_args(%arg2 = %2) -> (tensor<8xf32>) {
      %5 = affine.min #map(%arg1)
      %extracted_slice = tensor.extract_slice %arg0[0, %arg1] [8, %5] [1, 1] : tensor<1600x625xf32> to tensor<8x?xf32>
      %6 = vector.create_mask %c8, %5 : vector<8x8xi1>
      %7 = vector.mask %6 { vector.transfer_read %extracted_slice[%c0, %c0], %0 {in_bounds = [true, true]} : tensor<8x?xf32>, vector<8x8xf32> } : vector<8x8xi1> -> vector<8x8xf32>
      %8 = vector.transfer_read %arg2[%c0], %0 {in_bounds = [true]} : tensor<8xf32>, vector<8xf32>
      %9 = vector.mask %6 { vector.multi_reduction <add>, %7, %8 [1] : vector<8x8xf32> to vector<8xf32> } : vector<8x8xi1> -> vector<8xf32>
      %10 = vector.transfer_write %9, %arg2[%c0] {in_bounds = [true]} : vector<8xf32>, tensor<8xf32>
      scf.yield %10 : tensor<8xf32>
    }

[bjacob]

The above issue is fixed by llvm/llvm-project#189477 but somehow, after lowering to LLVM, this doesn't make any difference and the NaN bug isn't fixed.

[bjacob]

Further minimized reproducer:

specialized_32.zip

Differences: mimized shapes down to 64x4 instead of 1600x25, dropped the division-by-0, no NANs involved at all anymore. The bug is thus not anymore about presence of NANs in the output, it is that the results differ if you edit the testcase to change native_vector_size to 64.

The earlier vector masking bug seems to be legitimate, but isn't what is still causing incorrect results here even after applying the fix for that. Also, the reason why disabling vector masking "fixes" the bug is that that basically entirely disables vectorization, resulting in scalar code.

Whatever the bug is, it is a vectorization bug.

[bjacob]

Narrowed it down further to a specific operand of a linalg.generic reading wrong values.

repro2.zip

Unzip, cd into it, run ./repro2.sh

It will print:

EXEC @main
EXEC @main
outputs differ, good, you managed to reproduce the bug!

What it does: it specializes testcase_template.mlir with 2 different native_vector_size values, here 8 and 16. This should not make any difference to the outputs, so the bug is that it does.

Compared to earlier testcases, this new testcase tricks the compiler into actually just forwarding one of its outputs. I couldn't actually simplify the linalg.generic without losing the reproduction, but I used a cmpf against a 1e+9 constant (so the compare always fails) to select the input value that I want to forward, instead of the values computed earlier.

The input/output tensors are dumped in /tmp:

/tmp/repro2_input0.npy.txt
/tmp/repro2_input3.npy.txt
/tmp/repro2_input8.npy.txt
/tmp/repro2_8_output0.npy.txt
/tmp/repro2_16_output0.npy.txt

The input being forwarded is /tmp/repro2_input8.npy.txt. It has shape 16x4 and its contents are:

[[ 0.  1.  2.  3.]
 [ 4.  5.  6.  7.]
 [ 8.  9. 10. 11.]
 [12. 13. 14. 15.]
 [16. 17. 18. 19.]
 [20. 21. 22. 23.]
 [24. 25. 26. 27.]
 [28. 29. 30. 31.]
 [32. 33. 34. 35.]
 [36. 37. 38. 39.]
 [40. 41. 42. 43.]
 [44. 45. 46. 47.]
 [48. 49. 50. 51.]
 [52. 53. 54. 55.]
 [56. 57. 58. 59.]
 [60. 61. 62. 63.]]

When native_vector_size = 16, we see from /tmp/repro2_16_output0.npy.txt that the generic reads it as the following correct values given the broadcasting to shape 16x4x4 creating that new innermost dimension 4 that the values are broadcasted to:

[[[ 0.  0.  0.  0.]
  [ 1.  1.  1.  1.]
  [ 2.  2.  2.  2.]
  [ 3.  3.  3.  3.]]

 [[ 4.  4.  4.  4.]
  [ 5.  5.  5.  5.]
  [ 6.  6.  6.  6.]
  [ 7.  7.  7.  7.]]

 [[ 8.  8.  8.  8.]
  [ 9.  9.  9.  9.]
  [10. 10. 10. 10.]
  [11. 11. 11. 11.]]

...

 [[60. 60. 60. 60.]
  [61. 61. 61. 61.]
  [62. 62. 62. 62.]
  [63. 63. 63. 63.]]]

With native_vector_size = 8, we see from /tmp/repro2_8_output0.npy.txt that the generic reads incorrect values instead, where the native vector size of 8 is actually visible (it shouldn't be visible!):

[[[ 0.  4.  0.  4.]
  [ 1.  5.  1.  5.]
  [ 2.  6.  2.  6.]
  [ 3.  7.  3.  7.]]

 [[ 0.  4.  0.  4.]
  [ 1.  5.  1.  5.]
  [ 2.  6.  2.  6.]
  [ 3.  7.  3.  7.]]

 [[ 8. 12.  8. 12.]
  [ 9. 13.  9. 13.]
  [10. 14. 10. 14.]
  [11. 15. 11. 15.]]

 [[ 8. 12.  8. 12.]
  [ 9. 13.  9. 13.]
  [10. 14. 10. 14.]
  [11. 15. 11. 15.]]

 [[16. 20. 16. 20.]
  [17. 21. 17. 21.]
  [18. 22. 18. 22.]
  [19. 23. 19. 23.]]

...

 [[48. 52. 48. 52.]
  [49. 53. 49. 53.]
  [50. 54. 50. 54.]
  [51. 55. 51. 55.]]

 [[56. 60. 56. 60.]
  [57. 61. 57. 61.]
  [58. 62. 58. 62.]
  [59. 63. 59. 63.]]

 [[56. 60. 56. 60.]
  [57. 61. 57. 61.]
  [58. 62. 58. 62.]
  [59. 63. 59. 63.]]]

[bjacob]

Fixed by llvm/llvm-project#190230