mediapipe/mediapipe/gpu/metal.bzl

"""Experimental Skylark rules for Apple's Metal.

This creates a .metallib file containing compiled Metal shaders.
Note that the default behavior in Xcode is to put all metal shaders into a
single "default.metallib", which can be loaded using the method
newDefaultLibrary in MTLDevice. Meanwhile, the metal_library rule creates a
named .metallib, which can be loaded using newLibraryWithFile:error:.

Example:

  metal_library(
      name = "my_shaders",
      srcs = ["my_shaders.metal"],
      hdrs = ["my_shaders.h"],
  )

This produces a "my_shaders.metallib".

The metal_library target can be added to the deps attribute of an objc_library.
The dependent objc_library can then access the headers declared by the
metal_library, if any.

The metal_library target can also be added to the resources attribute as a
simple data file, but in that case any declared headers are not visible to
dependent objc_library rules.
"""

load("@build_bazel_apple_support//lib:apple_support.bzl", "apple_support")
load("@bazel_skylib//lib:dicts.bzl", "dicts")

# This load statement is overriding the visibility of the internal implementation of rules_apple.
# This rule will be migrated to rules_apple in the future, hence the override. Please do not use
# this import anywhere else.
load(
    "@build_bazel_rules_apple//apple/internal:resources.bzl",
    "resources",
)

def _metal_compiler_args(ctx, src, obj, minimum_os_version, copts, diagnostics, deps_dump):
    """Returns arguments for metal compiler."""
    apple_fragment = ctx.fragments.apple

    platform = apple_fragment.single_arch_platform

    if not minimum_os_version:
        minimum_os_version = ctx.attr._xcode_config[apple_common.XcodeVersionConfig].minimum_os_for_platform_type(
            platform.platform_type,
        )

    args = copts + [
        "-arch",
        "air64",  # TODO: choose based on target device/cpu/platform?
        "-emit-llvm",
        "-c",
        "-gline-tables-only",
        "-isysroot",
        apple_support.path_placeholders.sdkroot(),
        "-ffast-math",
        "-serialize-diagnostics",
        diagnostics.path,
        "-o",
        obj.path,
        "-mios-version-min=%s" % minimum_os_version,
        "",
        src.path,
        "-MMD",
        "-MT",
        "dependencies",
        "-MF",
        deps_dump.path,
    ]
    return args

def _metal_compiler_inputs(srcs, hdrs, deps = []):
    """Determines the list of inputs required for a compile action."""

    cc_infos = [dep[CcInfo] for dep in deps if CcInfo in dep]

    dep_headers = depset(transitive = [
        cc_info.compilation_context.headers
        for cc_info in cc_infos
    ])

    return depset(srcs + hdrs, transitive = [dep_headers])

def _metal_library_impl(ctx):
    """Implementation for metal_library Skylark rule."""

    # A unique path for rule's outputs.
    objs_outputs_path = "{}.objs/".format(ctx.label.name)

    output_objs = []
    for src in ctx.files.srcs:
        basename = src.basename
        obj = ctx.actions.declare_file(objs_outputs_path + basename + ".air")
        output_objs.append(obj)
        diagnostics = ctx.actions.declare_file(objs_outputs_path + basename + ".dia")
        deps_dump = ctx.actions.declare_file(objs_outputs_path + basename + ".dat")

        args = (["metal"] +
                _metal_compiler_args(ctx, src, obj, ctx.attr.minimum_os_version, ctx.attr.copts, diagnostics, deps_dump))

        apple_support.run(
            ctx,
            xcode_path_resolve_level = apple_support.xcode_path_resolve_level.args,
            inputs = _metal_compiler_inputs(ctx.files.srcs, ctx.files.hdrs, ctx.attr.deps),
            outputs = [obj, diagnostics, deps_dump],
            mnemonic = "MetalCompile",
            executable = "/usr/bin/xcrun",
            arguments = args,
            use_default_shell_env = False,
            progress_message = ("Compiling Metal shader %s" %
                                (basename)),
        )

    output_lib = ctx.actions.declare_file(ctx.label.name + ".metallib")
    args = [
        "metallib",
        "-split-module",
        "-o",
        output_lib.path,
    ] + [x.path for x in output_objs]

    apple_support.run(
        ctx,
        xcode_path_resolve_level = apple_support.xcode_path_resolve_level.args,
        inputs = output_objs,
        outputs = (output_lib,),
        mnemonic = "MetalLink",
        executable = "/usr/bin/xcrun",
        arguments = args,
        progress_message = (
            "Linking Metal library %s" % ctx.label.name
        ),
    )

    objc_provider = apple_common.new_objc_provider(
        providers = [x[apple_common.Objc] for x in ctx.attr.deps if apple_common.Objc in x],
    )

    cc_infos = [dep[CcInfo] for dep in ctx.attr.deps if CcInfo in dep]
    if ctx.files.hdrs:
        cc_infos.append(
            CcInfo(
                compilation_context = cc_common.create_compilation_context(
                    headers = depset([f for f in ctx.files.hdrs]),
                ),
            ),
        )

    return [
        DefaultInfo(
            files = depset([output_lib]),
        ),
        objc_provider,
        cc_common.merge_cc_infos(cc_infos = cc_infos),
        # Return the provider for the new bundling logic of rules_apple.
        resources.bucketize_typed([output_lib], "unprocessed"),
    ]

METAL_LIBRARY_ATTRS = dicts.add(apple_support.action_required_attrs(), {
    "srcs": attr.label_list(allow_files = [".metal"], allow_empty = False),
    "hdrs": attr.label_list(allow_files = [".h"]),
    "deps": attr.label_list(providers = [["objc", CcInfo], [apple_common.Objc, CcInfo]]),
    "copts": attr.string_list(),
    "minimum_os_version": attr.string(),
})

metal_library = rule(
    implementation = _metal_library_impl,
    attrs = METAL_LIBRARY_ATTRS,
    fragments = ["apple", "objc", "swift"],
    output_to_genfiles = True,
)
"""
Builds a Metal library.

Args:
  srcs: Metal shader sources.
  hdrs: Header files used by the shader sources.
  deps: objc_library targets whose headers should be visible to the shaders.

The header files declared in this rule are also visible to any objc_library
rules that have it as a dependency, so that constants and typedefs can be
shared between Metal and Objective-C code.
"""