parsers.bzl

# 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
#
#      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.

"""Common routines for working with XLScc parser in bazel files."""

load(
    "//transpiler:fhe_common.bzl",
    "executable_attr",
    "run_with_stem",
)

_XLSCC = "@com_google_xls//xls/contrib/xlscc:xlscc"
_GET_TOP_FUNC_FROM_PROTO = "@com_google_xls//xls/contrib/xlscc:get_top_func_from_proto"

XlsCcOutputInfo = provider(
    """The output of compiling C++ using XLScc.""",
    fields = {
        "ir": "XLS IR file generated by XLScc compiler",
        "metadata": "XLS IR protobuf by XLScc compiler",
        "metadata_entry": "Text file containing the entry point for the program",
        "library_name": "Library name; if empty, stem is used to derive names.",
        "stem": "Name stem derived from input source C++ file (e.g., 'myfile' from 'myfile.cc'.)",
        "hdrs": "Input C++ headers",
        "hdr_files": "Input header files",
    },
)

def _get_top_func(ctx, library_name, metadata_file):
    """Extract the name of the entry function from the XLS metadata file."""
    return run_with_stem(
        ctx,
        library_name,
        [metadata_file],
        ".entry",
        ctx.executable._get_top_func_from_proto,
        [metadata_file.path],
    )

def _get_cc_to_xls_ir_library_name(ctx):
    """Derive a stem from a file name (e.g., myfile.cc -- myfile)."""
    (library_name, _, _) = ctx.attr.src.label.name.rpartition(".cc")
    return library_name

def _build_xls_ir(ctx, library_name):
    """Build the XLS IR from a C++ source.

    Args:
      ctx: The Blaze context.
      library_name: The stem for the output file.

    Returns:
      A File containing the generated IR and one containing metadata about
      the translated function (signature, etc.).
    """
    ir_file = ctx.actions.declare_file("%s.ir" % library_name)
    metadata_file = ctx.actions.declare_file("%s_meta.proto" % library_name)
    defines = ""
    if ctx.attr.defines:
        defines = "--defines " + ",".join(ctx.attr.defines)

    rlimit = ""
    if ctx.attr.z3_rlimit:
        rlimit = "--z3_rlimit %d" % (ctx.attr.z3_rlimit,)

    ctx.actions.run_shell(
        inputs = [ctx.file.src] + ctx.files.hdrs,
        outputs = [ir_file, metadata_file],
        tools = [ctx.executable._xlscc],
        command = "%s %s --meta_out %s %s %s > %s" % (
            ctx.executable._xlscc.path,
            ctx.file.src.path,
            metadata_file.path,
            defines,
            rlimit,
            ir_file.path,
        ),
    )
    return (ir_file, metadata_file, _get_top_func(ctx, library_name, metadata_file))

def _cc_to_xls_ir_impl(ctx):
    stem = _get_cc_to_xls_ir_library_name(ctx)
    library_name = ctx.attr.library_name or stem
    ir_file, metadata_file, metadata_entry_file = _build_xls_ir(ctx, library_name)

    outputs = [
        ir_file,
        metadata_file,
        metadata_entry_file,
    ]

    return [
        DefaultInfo(files = depset(outputs)),
        XlsCcOutputInfo(
            ir = depset([ir_file]),
            metadata = depset([metadata_file]),
            metadata_entry = depset([metadata_entry_file]),
            library_name = library_name,
            stem = stem,
            hdrs = ctx.attr.hdrs,
            hdr_files = ctx.files.hdrs,
        ),
    ]

cc_to_xls_ir = rule(
    doc = """
      This rule uses XLScc to parse C++ code to XLS IR.  It emits the IR
      file, a protobuf-metadata file, a file containing the entry point.
      """,
    implementation = _cc_to_xls_ir_impl,
    attrs = {
        "src": attr.label(
            doc = "A single C++ source file to transpile.",
            allow_single_file = [".cc"],
        ),
        "hdrs": attr.label_list(
            doc = "Any headers necessary for conversion to XLS IR.",
            allow_files = [".h"],
        ),
        "library_name": attr.string(
            doc = """
            The name used for the output files (<library_name>.cc and <library_name>.h);
            If not specified, the default is derived from the basename of the source file.
            """,
        ),
        "defines": attr.string_list(
            doc = """
            A list of defines to pass to xlscc of the form "NAME=VALUE". The
            end user can set `defines=["FOO=BAR", "BAZ=QUXX"]` in an
            fhe_cc_library rule or similar to propagate defines through to
            xlscc.
            """,
        ),
        "z3_rlimit": attr.int(
            doc = """
            The computation limit to pass to xlscc, to instruct its Z3 solver
            on how much effort to spend trying to prove that loops unwrap. In
            some cases where the break logic within a loop is complicated, but
            the loop bounds are simple, setting and/or lowering this limit can
            improve compilation speed. However, some loops may not unroll if
            this limit is set too low.

            If set to zero, Z3 is given no limit and may run indefinitely.
            """,
            default = 10000,
        ),
        "_xlscc": executable_attr(_XLSCC),
        "_get_top_func_from_proto": attr.label(
            default = Label(_GET_TOP_FUNC_FROM_PROTO),
            executable = True,
            cfg = "exec",
        ),
    },
)