STT-tensorflow/third_party/mlir/linalggen.bzl

40 lines
1.3 KiB
Python

"""BUILD extensions for MLIR linalg generation."""
def genlinalg(name, linalggen, src, linalg_outs):
"""genlinalg() generates code from a tc spec file.
Args:
name: The name of the build rule for use in dependencies.
linalggen: The binary used to produce the output.
src: The tc spec file.
linalg_outs: A list of tuples (opts, out), where each opts is a string of
options passed to linalggen, and the out is the corresponding output file
produced.
"""
for (opts, out) in linalg_outs:
# All arguments to generate the output except output destination.
base_args = [
"$(location %s)" % linalggen,
"%s" % opts,
"$(location %s)" % src,
]
rule_suffix = "_".join(opts.replace("-", "_").replace("=", "_").split(" "))
# Rule to generate code using generated shell script.
native.genrule(
name = "%s_%s_genrule" % (name, rule_suffix),
srcs = [src],
outs = [out],
tools = [linalggen],
cmd = ("echo " + " ".join(base_args) + " -o $@; " + " ".join(base_args) + " -o $@"),
)
# List of opts that do not generate cc files.
hdrs = [f for (opts, f) in linalg_outs]
native.cc_library(
name = name,
hdrs = hdrs,
textual_hdrs = hdrs,
)