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?