Strange behavior of `convert_element_type` on main branch

[jakevdp]

Running from github HEAD:

import jax.numpy as jnp
import jax

jax.config.update('jax_enable_x64', True)

x = jnp.float64(2.718)
x_f16 = x.astype('float16')

print(x, x_f16)
# 2.718 -15790.0

version info:

jax.print_environment_info()
# jax:    0.3.18
# jaxlib: 0.3.15
# numpy:  1.23.2
# python: 3.8.2 (v3.8.2:7b3ab5921f, Feb 24 2020, 17:52:18)  [Clang 6.0 (clang-600.0.57)]
# jax.devices (1 total, 1 local): [CpuDevice(id=0)]
# process_count: 1

I can only reproduce this locally on my non-M1 macbook; I've not been able to reproduce in Colab or on linux.

[nicholasjng]

FWIW, this does not reproduce on my macbook, albeit with a (locally built) jaxlib 0.3.18, also at HEAD:

>>> import jax.numpy as jnp
>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> x = jnp.float64(2.718)
>>> x_f16 = x.astype("float16")
>>> print(x, x_f16)
2.718 2.719

>>> jax.print_environment_info()
jax:    0.3.18
jaxlib: 0.3.18
numpy:  1.23.3
python: 3.10.6 (main, Aug 30 2022, 04:58:14) [Clang 13.1.6 (clang-1316.0.21.2.5)]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1

[soraros]

It does reproduce on my MacBook, with jaxlib 0.3.15, and I think that's where the problem lies.

jax.print_environment_info()
# jax:    0.3.18
# jaxlib: 0.3.15
# numpy:  1.23.3
# python: 3.10.6 (main, Aug 30 2022, 05:12:36) [Clang 13.1.6 (clang-1316.0.21.2.5)]
# jax.devices (1 total, 1 local): [CpuDevice(id=0)]
# process_count: 1

[hawkinsp]

This does not reproduce on my M1 Macbook pro with either jaxlib 0.3.15 or jaxlib from head.

For those of you where this fails, are you using Intel or ARM macbooks? Can you try building jaxlib from head and see if it reproduces? I'm speculating this is probably already fixed with an up to date jaxlib.

[jakevdp]

Mine is an intel macbook

[soraros]

Mine is also an intel MacBook. And the problem persists with a locally built jaxlib from HEAD.

In [1]: import jax.numpy as jnp
   ...: import jax
   ...: 
   ...: jax.config.update('jax_enable_x64', True)
   ...: 
   ...: x = jnp.float64(2.718)
   ...: x_f16 = x.astype('float16')
   ...: 
   ...: print(x, x_f16)
   ...: # 2.718 -15790.0
2.718 -15790.0

In [2]: jax.print_environment_info()
jax:    0.3.18
jaxlib: 0.3.18
numpy:  1.23.3
python: 3.10.6 (main, Aug 30 2022, 05:12:36) [Clang 13.1.6 (clang-1316.0.21.2.5)]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1

[hawkinsp]

Could you share the details of what CPU you have? Sharing the output of sysctl -a | grep machdep.cpu should do it.

[soraros]

@hawkinsp Here you go.

``` machdep.cpu.mwait.linesize_min: 64 machdep.cpu.mwait.linesize_max: 64 machdep.cpu.mwait.extensions: 3 machdep.cpu.mwait.sub_Cstates: 286531872 machdep.cpu.thermal.sensor: 1 machdep.cpu.thermal.dynamic_acceleration: 1 machdep.cpu.thermal.invariant_APIC_timer: 1 machdep.cpu.thermal.thresholds: 2 machdep.cpu.thermal.ACNT_MCNT: 1 machdep.cpu.thermal.core_power_limits: 1 machdep.cpu.thermal.fine_grain_clock_mod: 1 machdep.cpu.thermal.package_thermal_intr: 1 machdep.cpu.thermal.hardware_feedback: 0 machdep.cpu.thermal.energy_policy: 1 machdep.cpu.xsave.extended_state: 31 832 1088 0 machdep.cpu.xsave.extended_state1: 15 832 256 0 machdep.cpu.arch_perf.version: 4 machdep.cpu.arch_perf.number: 4 machdep.cpu.arch_perf.width: 48 machdep.cpu.arch_perf.events_number: 7 machdep.cpu.arch_perf.events: 0 machdep.cpu.arch_perf.fixed_number: 3 machdep.cpu.arch_perf.fixed_width: 48 machdep.cpu.cache.linesize: 64 machdep.cpu.cache.L2_associativity: 4 machdep.cpu.cache.size: 256 machdep.cpu.tlb.inst.large: 8 machdep.cpu.tlb.data.small: 64 machdep.cpu.tlb.data.small_level1: 64 machdep.cpu.address_bits.physical: 39 machdep.cpu.address_bits.virtual: 48 machdep.cpu.tsc_ccc.numerator: 200 machdep.cpu.tsc_ccc.denominator: 2 machdep.cpu.max_basic: 22 machdep.cpu.max_ext: 2147483656 machdep.cpu.vendor: GenuineIntel machdep.cpu.brand_string: Intel(R) Core(TM) i9-9980HK CPU @ 2.40GHz machdep.cpu.family: 6 machdep.cpu.model: 158 machdep.cpu.extmodel: 9 machdep.cpu.extfamily: 0 machdep.cpu.stepping: 13 machdep.cpu.feature_bits: 9221959987971750911 machdep.cpu.leaf7_feature_bits: 43804591 1073741824 machdep.cpu.leaf7_feature_bits_edx: 3154120192 machdep.cpu.extfeature_bits: 1241984796928 machdep.cpu.signature: 591597 machdep.cpu.brand: 0 machdep.cpu.features: FPU VME DE PSE TSC MSR PAE MCE CX8 APIC SEP MTRR PGE MCA CMOV PAT PSE36 CLFSH DS ACPI MMX FXSR SSE SSE2 SS HTT TM PBE SSE3 PCLMULQDQ DTES64 MON DSCPL VMX EST TM2 SSSE3 FMA CX16 TPR PDCM SSE4.1 SSE4.2 x2APIC MOVBE POPCNT AES PCID XSAVE OSXSAVE SEGLIM64 TSCTMR AVX1.0 RDRAND F16C machdep.cpu.leaf7_features: RDWRFSGS TSC_THREAD_OFFSET SGX BMI1 AVX2 SMEP BMI2 ERMS INVPCID FPU_CSDS MPX RDSEED ADX SMAP CLFSOPT IPT SGXLC MDCLEAR IBRS STIBP L1DF ACAPMSR SSBD machdep.cpu.extfeatures: SYSCALL XD 1GBPAGE EM64T LAHF LZCNT PREFETCHW RDTSCP TSCI machdep.cpu.logical_per_package: 16 machdep.cpu.cores_per_package: 8 machdep.cpu.microcode_version: 240 machdep.cpu.processor_flag: 5 machdep.cpu.core_count: 8 machdep.cpu.thread_count: 16 ```

[hawkinsp]

Could someone who experiences this problem please run the reproduction with the environment variable XLA_FLAGS=--xla_dump_to=/tmp/somewhere and share a zip file or similar with the directory of files it produces?

[soraros]

@hawkinsp Here you go again. dump.zip

[hawkinsp]

My current best guess is that this may be related to the _Float16 ABI changing in LLVM.

This program in essence compiles into a single call to a builtin named __truncdfhf2. The only explanation I can think of is that there is a mismatch of calling conventions for Mac x86-64.

Notably this changed in LLVM recently: https://reviews.llvm.org/D131172

[d0k]

LLVM changed the x86-64 fp16 abi from passing the value in integer registers to passing it in floating point registers. This means that __truncdfhf2 that's coming from the system is now incompatible with LLVM in jaxlib.

If I remember correctly __truncdfhf2 is provided by the system on macOS, not XCode. This would mean we're stuck on the wrong ABI until Apple decides to change it. Which seems unlikely to happen outside of a major release. I'm curious what will happen when they hit this change when updating Clang in XCode though.

This isn't an issue for f32->f16 because there's been a hardware instruction for it since Haswell, but there's none for f64->f16 (it exists in AVX512, but no mac was released with it).

I think we can work around this by making simple_orc_jit bind __truncdfhf2 to the fallback version in runtime_fp16.cc, which I fixed to use the correct ABI. I have no Intel mac around to test that.