# Copyright 2021 The JAX Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from functools import partial import hashlib import os import random import sys import tempfile import unittest from unittest import mock, SkipTest import warnings from absl.testing import absltest from jax.experimental import PartitionSpec as P from jax.experimental.compilation_cache import compilation_cache as cc from jax.experimental.maps import xmap from jax.experimental.pjit import pjit import jax from jax import jit, lax, pmap from jax._src.util import prod import jax._src.test_util as jtu from jax._src.lib import xla_bridge from jax._src.lib import xla_client import numpy as np from jax.config import config from jax._src.config import (persistent_cache_min_compile_time_secs, raise_persistent_cache_errors) config.parse_flags_with_absl() FLAGS = config.FLAGS @jtu.with_config(jax_raise_persistent_cache_errors=True, jax_persistent_cache_min_compile_time_secs=0) class CompilationCacheTest(jtu.JaxTestCase): def setUp(self): super().setUp() supported_platforms = ["tpu"] if "--xla_gpu_enable_xla_runtime_executable=true" in os.environ.get("XLA_FLAGS", ""): supported_platforms.append("gpu") if "--xla_cpu_use_xla_runtime=true" in os.environ.get("XLA_FLAGS", ""): supported_platforms.append("cpu") if jtu.device_under_test() not in supported_platforms: raise SkipTest("serialize executable only works on " + ",".join(supported_platforms)) def tearDown(self): super().tearDown() cc._cache = None def test_compile_options(self): compile_options_not_filled = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) compile_options_filled = self.filled_compile_options() filled_hash1 = self.get_hashed_value(cc._hash_compile_options, compile_options_filled) filled_hash2 = self.get_hashed_value(cc._hash_compile_options, compile_options_filled) not_filled_hash3 = self.get_hashed_value(cc._hash_compile_options, compile_options_not_filled) self.assertEqual(filled_hash1, filled_hash2) self.assertNotEqual(filled_hash1, not_filled_hash3) def test_executable_build_options(self): compile_options_not_filled = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) compile_options_filled = self.filled_compile_options() filled_hash1 = self.get_hashed_value(cc._hash_executable_build_options, compile_options_filled.executable_build_options) filled_hash2 = self.get_hashed_value(cc._hash_executable_build_options, compile_options_filled.executable_build_options) not_filled_hash3 = self.get_hashed_value(cc._hash_executable_build_options, compile_options_not_filled.executable_build_options) self.assertEqual(filled_hash1, filled_hash2) self.assertNotEqual(filled_hash1, not_filled_hash3) def test_debug_options(self): compile_options = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) hash1 = self.get_hashed_value(cc._hash_debug_options, compile_options.executable_build_options.debug_options) hash2 = self.get_hashed_value(cc._hash_debug_options, compile_options.executable_build_options.debug_options) self.assertEqual(hash1, hash2) new_debug_options = self.create_new_debug_options(compile_options.executable_build_options.debug_options) hash3 = self.get_hashed_value(cc._hash_debug_options, new_debug_options) self.assertNotEqual(hash1, hash3) def test_hash_platform(self): hash1 = self.get_hashed_value(cc._hash_platform, xla_bridge.get_backend()) hash2 = self.get_hashed_value(cc._hash_platform, xla_bridge.get_backend()) self.assertEqual(hash1, hash2) if xla_bridge.get_backend().platform != "cpu": cpu_backend = xla_bridge.get_backend("cpu") hash3 = self.get_hashed_value(cc._hash_platform, cpu_backend) self.assertNotEqual(hash1, hash3) def test_hash_int(self): hash1 = self.get_hashed_value(cc._hash_int, 90) hash2 = self.get_hashed_value(cc._hash_int, 8) hash3 = self.get_hashed_value(cc._hash_int, 8) self.assertEqual(hash2, hash3) self.assertNotEqual(hash1, hash2) def test_hash_bool(self): hash1 = self.get_hashed_value(cc._hash_bool, False) hash2 = self.get_hashed_value(cc._hash_bool, True) hash3 = self.get_hashed_value(cc._hash_bool, True) self.assertEqual(hash2, hash3) self.assertNotEqual(hash1, hash2) def test_hash_string(self): hash1 = self.get_hashed_value(cc._hash_string, "foo") hash2 = self.get_hashed_value(cc._hash_string, "bar") hash3 = self.get_hashed_value(cc._hash_string, "bar") self.assertEqual(hash2, hash3) self.assertNotEqual(hash1, hash2) def test_same_hash_key(self): computation = jax.xla_computation(lambda x, y: x + y)(1, 1) compile_options = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) backend = xla_bridge.get_backend() self.assertEqual(cc.get_cache_key(computation, compile_options, backend), cc.get_cache_key(computation, compile_options, backend)) def test_different_hash_key(self): computation = jax.xla_computation(lambda x, y: x + y)(1, 1) compile_options_not_filled = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) compile_options_filled = self.filled_compile_options() backend = xla_bridge.get_backend() self.assertNotEqual(cc.get_cache_key(computation, compile_options_not_filled, backend), cc.get_cache_key(computation, compile_options_filled, backend)) def test_different_computations(self): computation1 = jax.xla_computation(lambda x, y: x + y)(1, 1) computation2 = jax.xla_computation(lambda x, y: x * y)(2, 2) compile_options = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) backend = xla_bridge.get_backend() self.assertNotEqual(cc.get_cache_key(computation1, compile_options, backend), cc.get_cache_key(computation2, compile_options, backend)) def test_xla_flags(self): if jtu.is_device_tpu_v4(): raise unittest.SkipTest("TODO(b/240151176)") computation = jax.xla_computation(lambda x, y: x + y)(1, 1) compile_options = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) backend = xla_bridge.get_backend() orig_xla_flags = os.getenv("XLA_FLAGS") orig_argv = sys.argv try: os.environ["XLA_FLAGS"] = "--xla_gpu_autotune_level=0" key1 = cc.get_cache_key(computation, compile_options, backend) os.environ["XLA_FLAGS"] = "--xla_gpu_autotune_level=1" key2 = cc.get_cache_key(computation, compile_options, backend) self.assertNotEqual(key1, key2) os.environ["XLA_FLAGS"] = "--xla_gpu_autotune_level=0" key3 = cc.get_cache_key(computation, compile_options, backend) self.assertEqual(key1, key3) # Test flag in _xla_flags_to_exclude_from_cache_key os.environ["XLA_FLAGS"] = ( "--xla_gpu_autotune_level=0 --xla_force_host_platform_device_count=8") key4 = cc.get_cache_key(computation, compile_options, backend) self.assertEqual(key1, key4) # Test flags given on command line del os.environ["XLA_FLAGS"] sys.argv.append("--xla_gpu_autotune_level=0") key5 = cc.get_cache_key(computation, compile_options, backend) self.assertEqual(key1, key5) sys.argv.append("--xla_force_host_platform_device_count=8") self.assertEqual(key1, key5) finally: if orig_xla_flags is not None: os.environ["XLA_FLAGS"] = orig_xla_flags elif os.getenv("XLA_FLAGS") is not None: del os.environ["XLA_FLAGS"] sys.argv = orig_argv def test_get_no_executable(self): with tempfile.TemporaryDirectory() as tmpdir: cc.initialize_cache(tmpdir) computation = jax.xla_computation(lambda x, y: x + y)(1, 1) compile_options = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) backend = xla_bridge.get_backend() self.assertEqual(cc.get_executable(computation, compile_options, backend), None) def test_diff_executables(self): with tempfile.TemporaryDirectory() as tmpdir: cc.initialize_cache(tmpdir) computation1 = jax.xla_computation(lambda x, y: x + y)(1, 1) computation2 = jax.xla_computation(lambda x, y: x * y)(2, 2) compile_options = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) backend = xla_bridge.get_backend() executable1 = backend.compile(computation1, compile_options) executable2 = backend.compile(computation2, compile_options) cc.put_executable("computation1", computation1, compile_options, executable1, backend) cc.put_executable("computation2", computation2, compile_options, executable2, backend) self.assertNotEqual(cc.get_executable(computation1, compile_options, backend), cc.get_executable(computation2, compile_options, backend)) def test_put_executable(self): with tempfile.TemporaryDirectory() as tmpdir: cc.initialize_cache(tmpdir) computation = jax.xla_computation(lambda x, y: x + y)(np.int32(1), np.int32(1)) compile_options = xla_bridge.get_compile_options( num_replicas=1, num_partitions=1) backend = xla_bridge.get_backend() executable = backend.compile(computation, compile_options) cc.put_executable("alambda", computation, compile_options, executable, backend) deserialized_executable = cc.get_executable(computation, compile_options, backend) inputs_to_executable = (np.array(1, dtype=np.int32), np.array(2, dtype=np.int32)) expected = xla_client.execute_with_python_values(executable, inputs_to_executable, backend) actual = xla_client.execute_with_python_values(deserialized_executable, inputs_to_executable, backend) self.assertEqual(expected, actual) def test_pmap(self): with tempfile.TemporaryDirectory() as tmpdir: cc.initialize_cache(tmpdir) f = pmap(lambda x: x - lax.psum(x, 'i'), axis_name='i') x = np.arange(jax.device_count(), dtype=np.int64) f(x) files_in_directory = len(os.listdir(tmpdir)) self.assertEqual(files_in_directory, 1) x = np.arange(jax.device_count(), dtype=np.float32) f(x) files_in_directory = len(os.listdir(tmpdir)) self.assertEqual(files_in_directory, 2) #TODO: create a test for calling pmap with the same input more than once def test_jit(self): with tempfile.TemporaryDirectory() as tmpdir: cc.initialize_cache(tmpdir) f = jit(lambda x: x*x) f(1) files_in_directory = len(os.listdir(tmpdir)) self.assertEqual(files_in_directory, 1) f(1.0) files_in_directory = len(os.listdir(tmpdir)) self.assertEqual(files_in_directory, 2) @jtu.with_mesh([('x', 2)]) def test_pjit(self): with tempfile.TemporaryDirectory() as tmpdir: cc.initialize_cache(tmpdir) @partial(pjit, in_axis_resources=(P('x'), P('x')), out_axis_resources=None) def f(x, y): return x + y shape = (8, 8) x = np.arange(prod(shape), dtype=np.int64).reshape(shape) f(x, x + 1) files_in_directory = len(os.listdir(tmpdir)) self.assertEqual(files_in_directory, 1) x = np.arange(prod(shape), dtype=np.float32).reshape(shape) f(x, x + 1) files_in_directory = len(os.listdir(tmpdir)) self.assertEqual(files_in_directory, 2) @jtu.with_mesh([('x', 2)]) def test_xmap(self): with tempfile.TemporaryDirectory() as tmpdir: cc.initialize_cache(tmpdir) def f(x): return x * 2 devices = np.array(jax.local_devices()[:2]) if devices.size < 2: raise SkipTest("Test requires 2 devices") x = np.arange(8, dtype=np.int64).reshape((2, 2, 2)) xmap(f, in_axes=['a', ...], out_axes=['a', ...], axis_resources={'a': 'x'})(x) files_in_directory = len(os.listdir(tmpdir)) self.assertEqual(files_in_directory, 1) x = np.arange(8, dtype=np.float32).reshape((2, 2, 2)) xmap(f, in_axes=['a', ...], out_axes=['a', ...], axis_resources={'a': 'x'})(x) files_in_directory = len(os.listdir(tmpdir)) self.assertEqual(files_in_directory, 2) def test_cache_write_warning(self): with tempfile.TemporaryDirectory() as tmpdir: cc.initialize_cache(tmpdir) f = jit(lambda x: x*x) with raise_persistent_cache_errors(False), \ mock.patch.object(cc._cache.__class__, 'put') as mock_put, \ warnings.catch_warnings(record=True) as w: mock_put.side_effect = RuntimeError("test error") self.assertEqual(f(2), 4) self.assertLen(w, 1) self.assertIn( "Error writing persistent compilation cache entry " "for 'jit__lambda_': RuntimeError: test error", str(w[0].message)) def test_cache_read_warning(self): with tempfile.TemporaryDirectory() as tmpdir: cc.initialize_cache(tmpdir) f = jit(lambda x: x*x) with raise_persistent_cache_errors(False), \ mock.patch.object(cc._cache.__class__, 'get') as mock_get, \ warnings.catch_warnings(record=True) as w: mock_get.side_effect = RuntimeError("test error") self.assertEqual(f(2), 4) self.assertLen(w, 1) self.assertIn( "Error reading persistent compilation cache entry " "for 'jit__lambda_': RuntimeError: test error", str(w[0].message)) def test_min_compile_time(self): with tempfile.TemporaryDirectory() as tmpdir, \ persistent_cache_min_compile_time_secs(2): cc.initialize_cache(tmpdir) # Mock time to progress in small intervals so compilation time is small. with mock.patch("time.monotonic", side_effect=np.arange(0, 10, .1)): jit(lambda x: x + 1)(1) files_in_cache = len(os.listdir(tmpdir)) self.assertEqual(files_in_cache, 0) # Mock time to progress in large intervals so compilation time is large. with mock.patch("time.monotonic", side_effect=np.arange(0, 100, 10)): jit(lambda x: x + 2)(1) files_in_cache = len(os.listdir(tmpdir)) self.assertEqual(files_in_cache, 1) def create_new_debug_options(self, debug_options_obj): debug_options_obj.xla_cpu_enable_fast_math = False debug_options_obj.xla_cpu_fast_math_honor_infs = False debug_options_obj.xla_cpu_fast_math_honor_nans = False debug_options_obj.xla_cpu_fast_math_honor_division = False debug_options_obj.xla_cpu_fast_math_honor_functions = False debug_options_obj.xla_gpu_enable_fast_min_max = False debug_options_obj.xla_backend_optimization_level = random.randint(0, 10) debug_options_obj.xla_cpu_enable_xprof_traceme = False debug_options_obj.xla_llvm_disable_expensive_passes = False debug_options_obj.xla_test_all_input_layouts = False return debug_options_obj def filled_compile_options(self): compile_options = xla_client.CompileOptions() compile_options.num_replicas = 1 compile_options.num_partitions = 1 shape = xla_client.Shape.array_shape(np.dtype(np.float32), [2]) shape_array = [shape, shape] compile_options.argument_layouts = shape_array compile_options.executable_build_options.result_layout = shape device_assignment = xla_client.DeviceAssignment.create(np.ndarray(shape=(2,2))) compile_options.device_assignment = device_assignment compile_options.executable_build_options.device_assignment = device_assignment return compile_options def get_hashed_value(self, hash_function, hash_function_input): hash_obj = hashlib.sha256() hash_function(hash_obj, hash_function_input) return hash_obj.digest().hex() if __name__ == "__main__": absltest.main(testLoader=jtu.JaxTestLoader())