JAX tests fail with "cpu" backend on aarch64

[sergey-kozub]

Running the following command fails on GB200:
$ pytest -sv tests/export_harnesses_multi_platform_test.py

With error message:

tests/export_harnesses_multi_platform_test.py::PrimitiveTest::test_prim_convert_element_type_dtypes_to_dtypes_shap[136/4822]
_100_100_olddtype_bfloat16_newdtype_int8 LLVM ERROR: Cannot select: 0xfac7241d7e90: nxv8i16 = AArch64ISD::FCVTZS_MERGE_PASST
HRU 0xfac7241d7e20, 0xfac724082ea0, undef:nxv8i16
  0xfac7241d7e20: nxv8i1 = AArch64ISD::PTRUE TargetConstant:i32<31>
    0xfac7241d7db0: i32 = TargetConstant<31>
  0xfac724082ea0: nxv8bf16 = vselect 0xfac7240829d0, 0xfac724082420, 0xfac724082c00
    0xfac7240829d0: nxv8i1 = setcc 0xfac724082c00, 0xfac724082420, setole:ch
      0xfac724082c00: nxv8bf16,ch = load<(invariant load (<vscale x 1 x s128>) from %ir.scevgep19, align 8, !noalias !4)> 0x
fac72416d8e0, 0xfac7240827a0, undef:i64
        0xfac7240827a0: i64 = add 0xfac7240825e0, 0xfac724082500
          0xfac7240825e0: i64,ch = CopyFromReg 0xfac72416d8e0, Register:i64 %7
            0xfac724082650: i64 = Register %7
          0xfac724082500: i64 = vscale Constant:i64<16>
            0xfac7240822d0: i64 = Constant<16>
        0xfac724082490: i64 = undef
      0xfac724082420: nxv8bf16 = splat_vector ConstantFP:bf16<APFloat(49920)>
        0xfac7240821f0: bf16 = ConstantFP<APFloat(49920)>
    0xfac724082420: nxv8bf16 = splat_vector ConstantFP:bf16<APFloat(49920)>
      0xfac7240821f0: bf16 = ConstantFP<APFloat(49920)>
    0xfac724082c00: nxv8bf16,ch = load<(invariant load (<vscale x 1 x s128>) from %ir.scevgep19, align 8, !noalias !4)> 0xfa
c72416d8e0, 0xfac7240827a0, undef:i64
      0xfac7240827a0: i64 = add 0xfac7240825e0, 0xfac724082500
        0xfac7240825e0: i64,ch = CopyFromReg 0xfac72416d8e0, Register:i64 %7
          0xfac724082650: i64 = Register %7
        0xfac724082500: i64 = vscale Constant:i64<16>
          0xfac7240822d0: i64 = Constant<16>
      0xfac724082490: i64 = undef
  0xfac7241d7790: nxv8i16 = undef
In function: convert.2_kernel
Fatal Python error: Aborted

"lscpu" output:

Architecture:             aarch64
  CPU op-mode(s):         64-bit
  Byte Order:             Little Endian
CPU(s):                   144
  On-line CPU(s) list:    0-143
Vendor ID:                ARM
  Model:                  0
  Thread(s) per core:     1
  Core(s) per socket:     72
  Socket(s):              2
  Stepping:               r0p0
  Frequency boost:        disabled
  CPU max MHz:            3429.0000
  CPU min MHz:            81.0000
  BogoMIPS:               2000.00
  Flags:                  fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm ssbs sb paca pacg dcpo
                          dp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh bti