tensorflow/tensorflow/compiler/xla/tests/build_defs.bzl

"""Build rules for XLA testing."""

load("@local_config_cuda//cuda:build_defs.bzl", "cuda_is_configured")
load("//tensorflow/compiler/xla/tests:plugin.bzl", "plugins")

all_backends = ["cpu", "cpu_parallel", "gpu"] + plugins.keys()

def filter_backends(backends):
  """Removes "gpu" from a backend list if CUDA is not enabled.

  This allows us to simply hardcode lists including "gpu" here and in the
  BUILD file, without causing failures when CUDA isn't enabled.'

  Args:
    backends: A list of backends to filter.

  Returns:
    The filtered list of backends.
  """
  if cuda_is_configured():
    return backends
  else:
    return [backend for backend in backends if backend != "gpu"]


def xla_test(name,
             srcs,
             deps,
             xla_test_library_deps=[],
             backends=[],
             args=[],
             tags=[],
             copts=[],
             data=[],
             backend_tags={},
             backend_args={},
             **kwargs):
  """Generates cc_test targets for the given XLA backends.

  This rule generates a cc_test target for one or more XLA backends and also
  a platform-agnostic cc_library rule. The arguments are identical to cc_test
  with two additions: 'backends' and 'backend_args'. 'backends' specifies the
  backends to generate tests for ("cpu", "cpu_parallel", "gpu"), and
  'backend_args'/'backend_tags' specifies backend-specific args parameters to
  use when generating the cc_test.

  The name of the cc_tests are the provided name argument with the backend name
  appended, and the cc_library target name is the provided name argument with
  "_lib" appended. For example, if name parameter is "foo_test", then the cpu
  test target will be "foo_test_cpu" and the cc_library target is "foo_lib".

  The cc_library target can be used to link with other plugins outside of
  xla_test.

  The build rule also defines a test suite ${name} which includes the tests for
  each of the supported backends.

  Each generated cc_test target has a tag indicating which backend the test is
  for. This tag is of the form "xla_${BACKEND}" (eg, "xla_cpu"). These
  tags can be used to gather tests for a particular backend into a test_suite.

  Examples:

    # Generates the targets: foo_test_cpu and foo_test_gpu.
    xla_test(
        name = "foo_test",
        srcs = ["foo_test.cc"],
        backends = ["cpu", "gpu"],
        deps = [...],
    )

    # Generates the targets: bar_test_cpu and bar_test_gpu. bar_test_cpu
    # includes the additional arg "--special_cpu_flag".
    xla_test(
        name = "bar_test",
        srcs = ["bar_test.cc"],
        backends = ["cpu", "gpu"],
        backend_args = {"cpu": ["--special_cpu_flag"]}
        deps = [...],
    )

  The build rule defines the preprocessor macro XLA_TEST_BACKEND_${BACKEND}
  to the value 1 where ${BACKEND} is the uppercase name of the backend.

  Args:
    name: Name of the target.
    srcs: Sources for the target.
    deps: Dependencies of the target.
    xla_test_library_deps: If set, the generated test targets will depend on the
      respective cc_libraries generated by the xla_test_library rule.
    backends: A list of backends to generate tests for. Supported
      values: "cpu", "cpu_parallel", "gpu". If this list is empty, the test will
      be generated for all supported backends.
    args: Test arguments for the target.
    tags: Tags for the target.
    backend_args: A dict mapping backend name to list of additional args to
      use for that target.
    backend_tags: A dict mapping backend name to list of additional tags to
      use for that target.
  """
  test_names = []
  if not backends:
    backends = all_backends

  native.cc_library(
      name="%s_lib" % name,
      srcs=srcs,
      copts=copts,
      testonly=True,
      deps=deps + ["//tensorflow/compiler/xla/tests:test_macros_header"],
  )

  for backend in filter_backends(backends):
    test_name = "%s_%s" % (name, backend)
    this_backend_tags = ["xla_%s" % backend]
    this_backend_copts = []
    this_backend_args = backend_args.get(backend, [])
    this_backend_data = []
    if backend == "cpu":
      backend_deps = ["//tensorflow/compiler/xla/service:cpu_plugin"]
      backend_deps += ["//tensorflow/compiler/xla/tests:test_macros_cpu"]
    elif backend == "cpu_parallel":
      backend_deps = ["//tensorflow/compiler/xla/service:cpu_plugin"]
      backend_deps += ["//tensorflow/compiler/xla/tests:test_macros_cpu"]
      this_backend_args += ["--xla_backend_extra_options=\"xla_cpu_parallel\""]
    elif backend == "gpu":
      backend_deps = ["//tensorflow/compiler/xla/service:gpu_plugin"]
      backend_deps += ["//tensorflow/compiler/xla/tests:test_macros_gpu"]
      this_backend_tags += ["requires-gpu-sm35"]
    elif backend in plugins:
      backend_deps = plugins[backend]["deps"]
      this_backend_copts += plugins[backend]["copts"]
      this_backend_tags += plugins[backend]["tags"]
      this_backend_args += plugins[backend]["args"]
      this_backend_data += plugins[backend]["data"]
    else:
      fail("Unknown backend %s" % backend)

    if xla_test_library_deps:
      for lib_dep in xla_test_library_deps:
        backend_deps += ["%s_%s" % (lib_dep, backend)]

    native.cc_test(
        name=test_name,
        srcs=srcs,
        tags=tags + backend_tags.get(backend, []) + this_backend_tags,
        copts=copts + ["-DXLA_TEST_BACKEND_%s=1" % backend.upper()] +
        this_backend_copts,
        args=args + this_backend_args,
        deps=deps + backend_deps,
        data=data + this_backend_data,
        **kwargs)

    test_names.append(test_name)

  native.test_suite(name=name, tests=test_names)

def xla_test_library(name,
                     srcs,
                     hdrs=[],
                     deps=[],
                     backends=[]):
  """Generates cc_library targets for the given XLA backends.

  This rule forces the sources to be compiled for each backend so that the
  backend specific macros could expand correctly. It's useful when test targets
  in different directories referring to the same sources but test with different
  arguments.

  Examples:

    # Generates the targets: foo_test_library_cpu and foo_test_gpu.
    xla_test_library(
        name = "foo_test_library",
        srcs = ["foo_test.cc"],
        backends = ["cpu", "gpu"],
        deps = [...],
    )
    # Then use the xla_test rule to generate test targets:
    xla_test(
        name = "foo_test",
        srcs = [],
        backends = ["cpu", "gpu"],
        deps = [...],
        xla_test_library_deps = [":foo_test_library"],
    )

  Args:
    name: Name of the target.
    srcs: Sources for the target.
    hdrs: Headers for the target.
    deps: Dependencies of the target.
    backends: A list of backends to generate libraries for.
      Supported values: "cpu", "cpu_parallel", "gpu". If this list is empty, the
      library will be generated for all supported backends.
  """

  if not backends:
    backends = all_backends

  for backend in filter_backends(backends):
    this_backend_copts = []
    if backend in ["cpu", "cpu_parallel", "gpu"]:
      backend_deps = ["//tensorflow/compiler/xla/tests:test_macros_%s" % backend]
    elif backend in plugins:
      backend_deps = plugins[backend]["deps"]
      this_backend_copts += plugins[backend]["copts"]
    else:
      fail("Unknown backend %s" % backend)

    native.cc_library(
        name = "%s_%s" % (name, backend),
        srcs = srcs,
        testonly = True,
        hdrs = hdrs,
        copts = ["-DXLA_TEST_BACKEND_%s=1" % backend.upper()]
        + this_backend_copts,
        deps = deps + backend_deps,
    )


def generate_backend_suites(backends=[]):
  if not backends:
    backends = all_backends
  for backend in filter_backends(backends):
    native.test_suite(name="%s_tests" % backend,
                      tags = ["xla_%s" % backend])


def generate_backend_test_macros(backends=[]):
  if not backends:
    backends = all_backends
  for backend in filter_backends(backends):
    manifest = ""
    if backend in plugins:
      manifest = plugins[backend]["disabled_manifest"]

    native.cc_library(
        name="test_macros_%s" % backend,
        testonly = True,
        srcs = ["test_macros.cc"],
        hdrs = ["test_macros.h"],
        copts = [
          "-DXLA_PLATFORM=\\\"%s\\\"" % backend.upper(),
          "-DXLA_DISABLED_MANIFEST=\\\"%s\\\"" % manifest,
        ],
        deps = [
            "//tensorflow/compiler/xla:types",
            "//tensorflow/core:lib",
            "//tensorflow/core:test",
        ])