Coverage for src/bartz/_workarounds.py: 94%

1# bartz/src/bartz/_workarounds.py 

2# 

3# Copyright (c) 2026, The Bartz Contributors 

4# 

5# This file is part of bartz. 

6# 

7# Permission is hereby granted, free of charge, to any person obtaining a copy 

8# of this software and associated documentation files (the "Software"), to deal 

9# in the Software without restriction, including without limitation the rights 

10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

11# copies of the Software, and to permit persons to whom the Software is 

12# furnished to do so, subject to the following conditions: 

13# 

14# The above copyright notice and this permission notice shall be included in all 

15# copies or substantial portions of the Software. 

16# 

17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

23# SOFTWARE. 

24 

25"""Workarounds for upstream bugs, applied when bartz is imported. 

26 

27Set the environment variable ``BARTZ_SKIP_XLA_WORKAROUND=1`` to disable them. 

28""" 

29 

30import os (empty)

31import re (empty)

32from importlib.util import find_spec (empty)

33 

34import jax (empty)

35 

36FTZ_ATOMICS_OPTION = '-nvptx-allow-ftz-atomics' (empty)

37 

38 

39def parse_version(version: str) -> tuple[int, int, int]: (empty)

40 """Extract the leading numeric triplet of a version string.""" 

41 match = re.match(r'^(\d+)\.(\d+)\.(\d+)', version) (empty)

42 assert match is not None, version (empty)

43 major, minor, patch = match.groups() (empty)

44 return int(major), int(minor), int(patch) (empty)

45 

46 

47def add_backend_extra_option(xla_flags: str, option: str) -> str: (empty)

48 """Return `xla_flags` with `option` added to ``--xla_backend_extra_options``.""" 

49 prefix = '--xla_backend_extra_options=' (empty)

50 tokens = xla_flags.split() (empty)

51 for i, token in enumerate(tokens): (empty)

52 if token.startswith(prefix): (empty)

53 value = token.removeprefix(prefix) (empty)

54 tokens[i] = prefix + (f'{value},{option}' if value else option) (empty)

55 break (empty)

56 else: 

57 tokens.append(prefix + option) (empty)

58 return ' '.join(tokens) (empty)

59 

60 

61def option_in_parsed_flags(option: str) -> bool: (empty)

62 """Check `option` is in the backend extra options XLA parsed from ``XLA_FLAGS``. 

63 

64 XLA reads ``XLA_FLAGS`` only once per process, on the first use of a jax 

65 backend; later changes to the environment variable are silently ignored. 

66 Constructing a `CompileOptions` triggers that one-time read, so calling 

67 this function right after modifying ``XLA_FLAGS`` both locks the change in 

68 (if XLA had not read the variable yet) and reports whether it was read. 

69 """ 

70 # import locally so that a future jaxlib dropping this internal module 

71 # can not break `import bartz` on jax versions that don't need the fix 

72 from jaxlib import xla_client # noqa: PLC0415 (empty)

73 

74 serialized = xla_client.CompileOptions().SerializeAsString() (empty)

75 return option.encode() in serialized (empty)

76 

77 

78def cuda_plugin_installed() -> bool: (empty)

79 """Check if a jax cuda plugin package is installed.""" 

80 return any(find_spec(f'jax_cuda{v}_plugin') is not None for v in (12, 13)) 

81 

82 

83def cuda_devices_available() -> bool: (empty)

84 """Check if jax can actually use cuda devices.""" 

85 try: (empty)

86 devices = jax.devices('cuda') (empty)

87 except RuntimeError: (empty)

88 return False (empty)

89 else: 

90 return bool(devices) (empty)

91 

92 

93def fix_gpu_scatter_performance() -> None: (empty)

94 """Restore native f32 atomics in gpu scatters on affected jax versions.""" 

95 # WORKAROUND(jax<=0.11.0): jax 0.10.2 and 0.11.0 ship an LLVM that lowers 

96 # the f32 atomic adds in gpu scatters to CAS loops, catastrophically slow 

97 # under index contention; this hidden LLVM option restores the native 

98 # atomics (and pre-0.10.2 numerics). See 

99 # https://github.com/jax-ml/jax/issues/38806. If the next jax release 

100 # contains the LLVM fix, delete this whole file and its uses. 

101 if not ((0, 10, 2) <= parse_version(jax.__version__) <= (0, 11, 0)): (empty)

102 return (empty)

103 if not cuda_plugin_installed(): 103 ↛ 104line 103 didn't jump to line 104 because the condition on line 103 was never true(empty)

104 return 

105 xla_flags = os.environ.get('XLA_FLAGS', '') (empty)

106 if FTZ_ATOMICS_OPTION in xla_flags or '--xla_gpu_ftz=true' in xla_flags.lower(): (empty)

107 return # the user already took care of it (empty)

108 os.environ['XLA_FLAGS'] = add_backend_extra_option(xla_flags, FTZ_ATOMICS_OPTION) (empty)

109 if not option_in_parsed_flags(FTZ_ATOMICS_OPTION) and cuda_devices_available(): 109 ↛ exitline 109 didn't return from function 'fix_gpu_scatter_performance' because the condition on line 109 was always true(empty)

110 msg = ( (empty)

111 f'jax {jax.__version__} has a severe performance regression in gpu ' 

112 'scatter operations (https://github.com/jax-ml/jax/issues/38806). ' 

113 f'bartz works around it by adding {FTZ_ATOMICS_OPTION} to the ' 

114 'XLA_FLAGS environment ' 

115 'variable, but XLA has already read XLA_FLAGS (that happens on the ' 

116 'first use of a jax backend), so the workaround is ineffective. ' 

117 'Either import bartz before using jax, or set ' 

118 f"XLA_FLAGS='--xla_backend_extra_options={FTZ_ATOMICS_OPTION}' " 

119 'before starting Python, or use a jax version other than ' 

120 '0.10.2/0.11.0. Set BARTZ_SKIP_XLA_WORKAROUND=1 to ignore this ' 

121 'error and run with slow gpu scatters.' 

122 ) 

123 raise RuntimeError(msg) (empty)

124 

125 

126def apply_workarounds() -> None: (empty)

127 """Apply all workarounds; invoked on ``import bartz``.""" 

128 if os.environ.get('BARTZ_SKIP_XLA_WORKAROUND', '') not in ('', '0'): (empty)

129 return (empty)

130 fix_gpu_scatter_performance() (empty)