multiarch.bzl

# Copyright 2017 Google Inc. All Rights Reserved.
#
# 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
#
#     http://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.
# =============================================================================

"""Build extension rules for handling multiple target architectures."""

# Build configs for specific CPU architectures. Each entry specified
# additional copts and tags.
# TODO(googleuser): Figure out a workaround for the shift
# instructions, and look for any other unsupported instructions.
MULTIARCH_CONFIGS = {
    "generic": {
        "copts": [],
        "tags": [],
    },
    "avx": {
        "copts": [
            "-msse4.2",
        ],
        "tags": [],
    },
    "avx2fma": {
        "copts": [
            "-msse4.2",
            "-mavx",
            "-mavx2",
            "-mfma",
        ],
        "tags": [
            "local",
            "manual",
        ],
    },
}

# List of targets which are built for multiple architectures. These
# dependencies in dragnn_cc_* build rules are replaced with one with the
# appropriate suffix, e.g. _multiarch_generic
MULTIARCH_TARGETS = [
    "//dragnn/runtime:biaffine_digraph_component",
    "//dragnn/runtime:bulk_dynamic_component",
    "//dragnn/runtime:bulk_feed_forward_network",
    "//dragnn/runtime:bulk_lstm_network",
    "//dragnn/runtime:feed_forward_network",
    "//dragnn/runtime:feed_forward_network_kernel",
    "//dragnn/runtime:feed_forward_network_layer",
    "//dragnn/runtime:fixed_embeddings",
    "//dragnn/runtime:linked_embeddings",
    "//dragnn/runtime:lstm_network",
    "//dragnn/runtime:lstm_network_kernel",
    "//dragnn/runtime:network_unit_base",
    "//dragnn/runtime:sequence_bulk_dynamic_component",
    "//dragnn/runtime:sequence_features",
    "//dragnn/runtime:sequence_links",
    "//dragnn/runtime:sequence_model",
    "//dragnn/runtime/lstm_cell:cell_function",
    "//dragnn/runtime/lstm_cell:test_helpers",
    "//dragnn/runtime/myelin:myelin_dynamic_component",
    "//dragnn/runtime/myelin:myelin_dynamic_component_base",
    "//dragnn/runtime/myelin:sequence_myelin_dynamic_component",
    "//dragnn/runtime/xla:sequence_xla_dynamic_component_mixin",
    "//dragnn/runtime/xla:testdata_simple_component_library",
    "//dragnn/runtime/xla:xla_aot_dynamic_component",
    "//dragnn/runtime/xla:xla_dynamic_component",
    "//dragnn/runtime/xla:xla_dynamic_component_base",
]

def multiarch_name(target_name, arch_name):
  """Generates the multiarch version of |target_name| given |arch_name|."""
  return target_name + '_multiarch_' + arch_name

def _is_multiarch(target):
  """Returns true if |target| is designated as a multiarch target."""
  return (target in MULTIARCH_TARGETS or
          ('//' + native.package_name() + target) in MULTIARCH_TARGETS)

def _dragnn_cc_multiarch_target(native_rule = None,
                                name = '',
                                target_arch = None,
                                target_suffix = '',
                                copts = [],
                                deps = [],
                                tags = [],
                                opts_self = False,
                                deps_transformer = None,
                                **kwargs):
  """Generates a target for multiple architectures.

  Using the |native_rule| (e.g. cc_library) to create a set of targets for
  all CPU architectures listed in MULTIARCH_CONFIGS, with added suffixes
  that designate the architecture.

  When |target_arch| is set, then only that single target is generated,
  and the name of the target is unchanged (no suffix is added).

  When |opts_self| is true, then the 'copts' entry in MULTIARCH_CONFIGS
  is additionally used to build this target.

  The 'tags' entry in MULTIARCH_CONFIGS are included in the build tags.

  Args:
    native_rule: The build rule used for all generated targets
    name: The original name of the build rule (without any suffix).
    target_arch: When set, only this architecture is targeted.
    target_suffix: Additional suffix to add after the architecture.
    copts: The original compilation options for this target.
    deps: The original dependencies for this target.
    tags: The original build tags for this target.
    opts_self: When true, additional copts are included.
    deps_transformer: When set, a function to apply to the multiarch deps.
    **kwargs: Additional args passed along to the build rule.
  """
  # Determine set of target architectures based on |target_arch|.
  if target_arch:
    if target_arch in MULTIARCH_CONFIGS:
      arch_items = [(target_arch, MULTIARCH_CONFIGS[target_arch])]
    else:
      fail('Unknown target_arch value: ' + target_arch)
  else:
    arch_items = MULTIARCH_CONFIGS.items()

  # There is one target for each architecture in |arch_items|.
  for arch, arch_config in arch_items:
    # Transform the multi-arch deps.
    multiarch_deps = [multiarch_name(dep, arch) if _is_multiarch(dep) else dep
                      for dep in deps]
    if deps_transformer:
      multiarch_deps = deps_transformer(multiarch_deps)

    native_rule(
        name = (name if target_arch else multiarch_name(name, arch)) + target_suffix,
        copts = copts + arch_config['copts'] if opts_self else copts,
        deps = multiarch_deps,
        tags = tags + arch_config['tags'],
        **kwargs)

def _dragnn_cc_multiarch_test_target(name = None,
                                     target_arch = None,
                                     **kwargs):
  """Test target wrapper which puts arch name before '_test'."""
  test_suffix = '_test'
  has_test_suffix = name.endswith(test_suffix)

  # Keeps _test at the end of the target name.
  test_name = name[:-len(test_suffix)] if has_test_suffix else name
  target_suffix = test_suffix if has_test_suffix else ''

  _dragnn_cc_multiarch_target(native_rule = native.cc_test,
                              name = test_name,
                              target_arch = target_arch,
                              target_suffix = target_suffix,
                              **kwargs)

  # When |target_arch| is set, the resulting test is named |name|.  Otherwise,
  # tests with arch-specific names are generated, and for convenience we add a
  # test_suite named |name| that runs the generic version of the test.
  if not target_arch:
    native.test_suite(
        name = name,
        tests = [multiarch_name(test_name, 'generic') + target_suffix])

def dragnn_cc_multiarch_library(**kwargs):
  """Similar to cc_library, but creates multiple architecture targets."""

  _dragnn_cc_multiarch_target(native_rule = native.cc_library,
                              **kwargs)

def dragnn_cc_multiarch_test(**kwargs):
  """Similar to cc_test, but creates multiple architecture targets."""
  _dragnn_cc_multiarch_test_target(**kwargs)

def dragnn_cc_multiarch_binary(**kwargs):
  """Similar to cc_binary, but creates multiple architecture targets."""
  _dragnn_cc_multiarch_target(native_rule = native.cc_binary,
                              **kwargs)

def dragnn_cc_library(target_arch = 'generic', **kwargs):
  """Similar to cc_library, but targets one specific architecture."""
  _dragnn_cc_multiarch_target(native_rule = native.cc_library,
                              target_arch = target_arch,
                              **kwargs)

def dragnn_cc_test(target_arch = 'generic', **kwargs):
  """Similar to cc_test, but targets one specific architecture."""
  _dragnn_cc_multiarch_test_target(target_arch = target_arch,
                                   **kwargs)

def dragnn_cc_binary(target_arch = 'generic', **kwargs):
  """Similar to cc_binary, but targets one specific architecture."""
  _dragnn_cc_multiarch_target(native_rule = native.cc_binary,
                              target_arch = target_arch,
                              **kwargs)