vector.shape_cast operation fails on s390x

[christopherlmunoz]

8 onnx-mlir backend inferencing tests fail blocking #3053. I am commenting out the failing tests to move along with LLVM & StableHLO update. I will continue the investigation here.

Example of failure.

/usr/lib64/python3.9/subprocess.py:528: CalledProcessError
----------------------------- Captured stdout call -----------------------------
[1/6] Thu Jan 30 08:57:35 2025 (0s) Importing ONNX Model to MLIR Module from "test_dynamicquantizelinear_max_adjusted_expanded.onnx"
[2/6] Thu Jan 30 08:57:36 2025 (1s) Compiling and Optimizing MLIR Module
----------------------------- Captured stderr call -----------------------------
failed to legalize operation 'vector.shape_cast'
_____ OnnxBackendNodeModelTest.test_dynamicquantizelinear_max_adjusted_cpu _____
[gw0] linux -- Python 3.9.20 /usr/bin/python3.9

Failing tests:
test_dynamicquantizelinear_max_adjusted_expanded
test_dynamicquantizelinear_max_adjusted
test_dynamicquantizelinear_expanded
test_quantizelinear
test_round
test_dynamicquantizelinear
test_dynamicquantizelinear_min_adjusted
test_dynamicquantizelinear_min_adjusted_expanded

[christopherlmunoz]

Commands used to generate failure:

Can generate failures in debug and release mode. I am testing on a z16 LoZ s390x environment.

cmake -DMLIR_DIR=$MLIR_DIR -DCMAKE_BUILD_TYPE=Release -DONNX_MLIR_TEST_OPTLEVEL=0 -DONNX_MLIR_ACCELERATORS=NNPA ..

make TEST_NOFLOAT16=true TEST_MARCH=arch14 TEST_ARGS="-march=arch14" check-onnx-backend-numerical -j8
make TEST_NOFLOAT16=true TEST_MARCH=arch14 TEST_OPTLEVEL=3 check-onnx-backend-numerical -j8

When TEST_OPTLEVEL is not set to 03 the failures do not occur.

Dockerfile I believe defaults to -03.

subprocess.CalledProcessError: Command '['/workdir/onnx-mlir/build/Release/bin/onnx-mlir', '-O3', '--march=z16', '--constants-to-file-total-threshold=2', '--EmitJNI', '/workdir/onnx-mlir/build/test/backend/Release/check-onnx-backend-jni/test_round/test_round.onnx', '-o=/workdir/onnx-mlir/build/test/backend/Release/check-onnx-backend-jni/test_round/test_round']' returned non-zero exit status 13.

/usr/lib64/python3.9/subprocess.py:528: CalledProcessError
----------------------------- Captured stdout call -----------------------------
[1/8] Thu Jan 30 08:58:14 2025 (0s) Importing ONNX Model to MLIR Module from "test_round.onnx"
[2/8] Thu Jan 30 08:58:14 2025 (0s) Compiling and Optimizing MLIR Module
----------------------------- Captured stderr call -----------------------------
failed to legalize operation 'vector.shape_cast'```

Failure:

root@5ce6f0a5f725:~/onnx-mlir/build# ./Debug/bin/onnx-mlir -O3 -march=arch14 -profile-ir=Onnx ../third_party/onnx/onnx/backend/test/data/node/test_round/model.onnx --opt-report=Simd
[1/6] Mon Feb 10 15:36:25 2025 (0s) Importing ONNX Model to MLIR Module from "model.onnx"
[2/6] Mon Feb 10 15:36:26 2025 (1s) Compiling and Optimizing MLIR Module
==SIMD-REPORT==, onnx.Round-simd, onnx.Round_0, unary fully flattened, 16, 15
failed to legalize operation 'vector.shape_cast'

Sucess:

root@5ce6f0a5f725:~/onnx-mlir/build# ./Debug/bin/onnx-mlir -O1 -march=arch14 -profile-ir=Onnx ../third_party/onnx/onnx/backend/test/data/node/test_round/model.onnx --opt-report=Simd
[1/6] Mon Feb 10 15:36:31 2025 (0s) Importing ONNX Model to MLIR Module from "model.onnx"
[2/6] Mon Feb 10 15:36:31 2025 (0s) Compiling and Optimizing MLIR Module
==SIMD-REPORT==, onnx.Round, onnx.Round_0, no simd in unary because scalar/layouts, 0, 0
[3/6] Mon Feb 10 15:36:31 2025 (0s) Translating MLIR Module to LLVM and Generating LLVM Optimized Bitcode
[4/6] Mon Feb 10 15:36:31 2025 (0s) Generating Object from LLVM Bitcode
[5/6] Mon Feb 10 15:36:31 2025 (0s) Linking and Generating the Output Shared Library
[6/6] Mon Feb 10 15:36:31 2025 (0s) Compilation completed

[christopherlmunoz]

Some other failures attempting to debug with onnx-mlir-opt

./Release/bin/onnx-mlir -O3 -maccel=NNPA -march=arch14 --EmitMLIR ../third_party/onnx/onnx/backend/test/data/node/test_round/model.onnx

Output:

root@5ce6f0a5f725:~/onnx-mlir/build# cat ../third_party/onnx/onnx/backend/test/data/node/test_round/model.onnx.mlir
module attributes {llvm.data_layout = "E-m:e-i1:8:16-i8:8:16-i64:64-f128:64-v128:64-a:8:16-n32:64", llvm.target_triple = "s390x-unknown-linux-gnu", "onnx-mlir.accels" = ["NNPA-0x10001"], "onnx-mlir.symbol-postfix" = "model"} {
  func.func @main_graph(%arg0: memref<15xf32> {onnx.name = "x"}) -> (memref<15xf32> {onnx.name = "y"}) attributes {llvm.emit_c_interface} {
    %c0 = arith.constant 0 : index
    %alloc = memref.alloc() {alignment = 16 : i64} : memref<64xi8>
    %view = memref.view %alloc[%c0][] : memref<64xi8> to memref<15xf32>
    affine.for %arg1 = 0 to 15 step 16 {
      %0 = vector.load %arg0[%arg1] : memref<15xf32>, vector<16xf32>
      %1 = vector.shape_cast %0 : vector<16xf32> to vector<4x4xf32>
      %2 = vector.extract %1[0] : vector<4xf32> from vector<4x4xf32>
      %3 = "krnl.round_even"(%2) : (vector<4xf32>) -> vector<4xf32>
      %4 = vector.insert %3, %1 [0] : vector<4xf32> into vector<4x4xf32>
      %5 = vector.extract %1[1] : vector<4xf32> from vector<4x4xf32>
      %6 = "krnl.round_even"(%5) : (vector<4xf32>) -> vector<4xf32>
      %7 = vector.insert %6, %4 [1] : vector<4xf32> into vector<4x4xf32>
      %8 = vector.extract %1[2] : vector<4xf32> from vector<4x4xf32>
      %9 = "krnl.round_even"(%8) : (vector<4xf32>) -> vector<4xf32>
      %10 = vector.insert %9, %7 [2] : vector<4xf32> into vector<4x4xf32>
      %11 = vector.extract %1[3] : vector<4xf32> from vector<4x4xf32>
      %12 = "krnl.round_even"(%11) : (vector<4xf32>) -> vector<4xf32>
      %13 = vector.insert %12, %10 [3] : vector<4xf32> into vector<4x4xf32>
      %14 = vector.shape_cast %13 : vector<4x4xf32> to vector<16xf32>
      vector.store %14, %view[%arg1] : memref<15xf32>, vector<16xf32>
    }
    return %view : memref<15xf32>
  }
  "krnl.entry_point"() {func = @main_graph, numInputs = 1 : i32, numOutputs = 1 : i32, signature = "[    { \22type\22 : \22f32\22 , \22dims\22 : [15] , \22name\22 : \22x\22 }\0A\0A]\00@[   { \22type\22 : \22f32\22 , \22dims\22 : [15] , \22name\22 : \22y\22 }\0A\0A]\00"} : () -> ()
}

root@5ce6f0a5f725:~/onnx-mlir/build# ./Release/bin/onnx-mlir-opt -O3 -maccel=NNPA -march=arch14 --rewrite-onnx-for-zhigh ../third_party/onnx/onnx/backend/test/data/node/test_round/model.onnx.mlir --convert-onnx-to-zhigh --convert-onnx-to-krnl --convert-krnl-to-affine
../third_party/onnx/onnx/backend/test/data/node/test_round/model.onnx.mlir:4:14: error: failed to legalize operation 'memref.alloc' that was explicitly marked illegal
    %alloc = memref.alloc() {alignment = 16 : i64} : memref<64xi8>
             ^
../third_party/onnx/onnx/backend/test/data/node/test_round/model.onnx.mlir:4:14: note: see current operation: %1 = "memref.alloc"() <{alignment = 16 : i64, operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<64xi8>

[christopherlmunoz]

@AlexandreEichenberger @tungld, please let me know if you can continue debugging this information or if you have any other questions. Are we at a point where we can go ahead with #3053 and try to resolve this ASAP?

[AlexandreEichenberger]

@christopherlmunoz will look into this

[AlexandreEichenberger]

@christopherlmunoz thanks for the instructions, I was able to reproduce the error with the Debug compiler, after uncommenting test_round and using this command

ONNX_MLIR_FLAGS="-O3 -march=arch15" TEST_CASE_BY_USER=test_round_cpu make -j 21 check-onnx-backend

It needs both flags to generate SIMD code, so it make sense. Let me now dive in

[AlexandreEichenberger]

FYI, this can be debugged also under x86/arm,... as we can use onnx-mlir with the target -O3 -march=arch14 on any platforms. Eventually it dies as we don't cross compile, but it goes down to llvm inside onnx-mlir.

I was able to reproduce that error on my Mac and solve it there, in case your environment works better on your laptop.

Also, if the fix that I suggested works (replacing the code in the if by the else):

#if 0
  vector::populateVectorToVectorCanonicalizationPatterns(patterns);
  vector::populateVectorBroadcastLoweringPatterns(patterns);
  vector::populateVectorContractLoweringPatterns(
      patterns, vector::VectorTransformsOptions());
  vector::populateVectorTransposeLoweringPatterns(
      patterns, vector::VectorTransformsOptions());
  vector::populateVectorShapeCastLoweringPatterns(patterns);
#else
  vector::populateVectorToVectorCanonicalizationPatterns(patterns);
  vector::populateVectorBitCastLoweringPatterns(patterns);
  vector::populateVectorBroadcastLoweringPatterns(patterns);
  vector::populateVectorContractLoweringPatterns(
      patterns, vector::VectorTransformsOptions());
  vector::populateVectorMaskOpLoweringPatterns(patterns);
  vector::populateVectorShapeCastLoweringPatterns(patterns);
  vector::populateVectorInterleaveLoweringPatterns(patterns);
  vector::populateVectorTransposeLoweringPatterns(
      patterns, vector::VectorTransformsOptions());
  // Vector transfer ops with rank > 1 should be lowered with VectorToSCF.
  vector::populateVectorTransferLoweringPatterns(
      patterns, /*maxTransferRank=*/1);
  vector::populateVectorMaskMaterializationPatterns(
      patterns, /*force32BitVectorIndices*/false);
  vector::populateVectorInsertExtractStridedSliceTransforms(patterns);
  vector::populateVectorStepLoweringPatterns(patterns);
  vector::populateVectorRankReducingFMAPattern(patterns);
#endif

This is something that we see from time to time when upgrading LLVM. First step is to see if they deprecated the op that does not lower. If it's still there, it's probable that they moved the patterns to a new "populate". Fortunately, there are other places in llvm-project/mlir where they use the patterns fully. So find one of the populate function we use, see where its used in mlir, and see if there are new calls that we don't currently call.