1"""BUILD extensions for MLIR linalg generation."""
2
3def genlinalg(name, linalggen, src, linalg_outs):
4    """genlinalg() generates code from a tc spec file.
5
6    Args:
7      name: The name of the build rule for use in dependencies.
8      linalggen: The binary used to produce the output.
9      src: The tc spec file.
10      linalg_outs: A list of tuples (opts, out), where each opts is a string of
11        options passed to linalggen, and the out is the corresponding output file
12        produced.
13    """
14
15    for (opts, out) in linalg_outs:
16        # All arguments to generate the output except output destination.
17        base_args = [
18            "$(location %s)" % linalggen,
19            "%s" % opts,
20            "$(location %s)" % src,
21        ]
22        rule_suffix = "_".join(opts.replace("-", "_").replace("=", "_").split(" "))
23
24        # Rule to generate code using generated shell script.
25        native.genrule(
26            name = "%s_%s_genrule" % (name, rule_suffix),
27            srcs = [src],
28            outs = [out],
29            tools = [linalggen],
30            cmd = (" ".join(base_args) + " -o $@"),
31        )
32
33    # List of opts that do not generate cc files.
34    hdrs = [f for (opts, f) in linalg_outs]
35    native.cc_library(
36        name = name,
37        hdrs = hdrs,
38        textual_hdrs = hdrs,
39    )
40