compilation_cache_test.py

# Copyright 2021 Google LLC
#
# 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 SkipTest

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
import jax._src.lib
import numpy as np

from jax.config import config
config.parse_flags_with_absl()
FLAGS = config.FLAGS

class CompilationCacheTest(jtu.JaxTestCase):

  def setUp(self):
    super().setUp()
    if jtu.device_under_test() != "tpu":
        raise SkipTest("serialize executable only works on TPU")
    if jax._src.lib.xla_bridge.get_backend().runtime_type == "tfrt":
        raise SkipTest("the new TFRT runtime does not support serialization")

  def tearDown(self):
      super().tearDown()
      cc._cache = None

  @unittest.skipIf(jax._src.lib.version < (0, 1, 68), "fails with earlier jaxlibs")
  def test_compile_options(self):
    compile_options_not_filled = jax._src.lib.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 = jax._src.lib.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 = jax._src.lib.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, jax._src.lib.xla_bridge.get_backend())
    hash2 = self.get_hashed_value(cc._hash_platform, jax._src.lib.xla_bridge.get_backend())
    self.assertEqual(hash1, hash2)
    if jax._src.lib.xla_bridge.get_backend().platform != "cpu":
        cpu_backend = jax._src.lib.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 = jax._src.lib.xla_bridge.get_compile_options(
                       num_replicas=1, num_partitions=1)
    backend = jax._src.lib.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 = jax._src.lib.xla_bridge.get_compile_options(
                       num_replicas=1, num_partitions=1)
    compile_options_filled = self.filled_compile_options()
    backend = jax._src.lib.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 = jax._src.lib.xla_bridge.get_compile_options(
                       num_replicas=1, num_partitions=1)
    backend = jax._src.lib.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):
    computation = jax.xla_computation(lambda x, y: x + y)(1, 1)
    compile_options = jax._src.lib.xla_bridge.get_compile_options(
        num_replicas=1, num_partitions=1)
    backend = jax._src.lib.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 = jax._src.lib.xla_bridge.get_compile_options(
          num_replicas=1, num_partitions=1)
      backend = jax._src.lib.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 = jax._src.lib.xla_bridge.get_compile_options(
          num_replicas=1, num_partitions=1)
      backend = jax._src.lib.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)(1, 1)
      compile_options = jax._src.lib.xla_bridge.get_compile_options(
          num_replicas=1, num_partitions=1)
      backend = jax._src.lib.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 = jax._src.lib.xla_client.execute_with_python_values(executable, inputs_to_executable, backend)
      actual = jax._src.lib.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 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 = jax._src.lib.xla_client.CompileOptions()
    compile_options.num_replicas = 1
    compile_options.num_partitions = 1
    shape = jax._src.lib.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 = jax._src.lib.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())