utils/bazel/llvm-project-overlay/mlir/tblgen.bzl

   1 # This file is licensed under the Apache License v2.0 with LLVM Exceptions.
   2 # See https://llvm.org/LICENSE.txt for license information.
   3 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   4 """BUILD extensions for MLIR table generation."""
   5
   6 load("@bazel_skylib//lib:paths.bzl", "paths")
   7
   8 TdInfo = provider(
   9     "Holds TableGen files and the dependencies and include paths necessary to" +
  10     " build them.",
  11     fields = {
  12         "transitive_sources": "td files transitively used by this rule.",
  13         "transitive_includes": (
  14             "include arguments to add to the final TableGen invocation. These" +
  15             " are the absolute directory paths that will be added with '-I'."
  16         ),
  17     },
  18 )
  19
  20 # For now we allow anything that provides DefaultInfo to just forward its files.
  21 # In particular, this allows filegroups to be used. This is mostly to ease
  22 # transition. In the future, the TdInfo provider will be required.
  23 # TODO(gcmn): Switch to enforcing TdInfo provider.
  24 def _get_dep_transitive_srcs(dep):
  25     """Extract TdInfo.transitive_sources, falling back to DefaultInfo.files."""
  26     if TdInfo in dep:
  27         return dep[TdInfo].transitive_sources
  28     return dep[DefaultInfo].files
  29
  30 def _get_dep_transitive_includes(dep):
  31     """Extract TdInfo.transitive_includes, falling back to an empty depset()."""
  32     if TdInfo in dep:
  33         return dep[TdInfo].transitive_includes
  34     return depset()
  35
  36 def _get_transitive_srcs(srcs, deps):
  37     """Obtain the source files for a target and its transitive dependencies.
  38
  39     Args:
  40       srcs: a list of source files
  41       deps: a list of targets that are direct dependencies
  42     Returns:
  43       a collection of the transitive sources
  44     """
  45     return depset(
  46         direct = srcs,
  47         transitive = [_get_dep_transitive_srcs(dep) for dep in deps],
  48     )
  49
  50 def _get_transitive_includes(includes, deps):
  51     """Obtain the includes paths for a target and its transitive dependencies.
  52
  53     Args:
  54       includes: a list of include paths
  55       deps: a list of targets that are direct dependencies
  56     Returns:
  57       a collection of the transitive include paths
  58     """
  59     return depset(
  60         direct = includes,
  61         transitive = [_get_dep_transitive_includes(dep) for dep in deps],
  62     )
  63
  64 def _prefix_roots(ctx, includes):
  65     """Map the given includes to be relative to all root directories.
  66
  67     This will expand them to be relative to all the root directories available
  68     in the execution environment for ctx.run (bin and genfiles in addition to
  69     the normal source root)
  70     """
  71     prefixed_includes = []
  72     for include in includes:
  73         prefixed_includes.append(include)
  74         prefixed_includes.append(paths.join(ctx.genfiles_dir.path, include))
  75         prefixed_includes.append(paths.join(ctx.bin_dir.path, include))
  76     return prefixed_includes
  77
  78 def _resolve_includes(ctx, includes):
  79     """Resolves include paths to paths relative to the execution root.
  80
  81     Relative paths are interpreted as relative to the current label's package.
  82     Absolute paths are interpreted as relative to the current label's workspace
  83     root."""
  84     package = ctx.label.package
  85     workspace_root = ctx.label.workspace_root
  86     workspace_root = workspace_root if workspace_root else "."
  87     resolved_includes = []
  88     for include in includes:
  89         if paths.is_absolute(include):
  90             include = include.lstrip("/")
  91         else:
  92             include = paths.join(package, include)
  93         include = paths.join(workspace_root, include)
  94         resolved_includes.extend(_prefix_roots(ctx, [include]))
  95     return resolved_includes
  96
  97 def _td_library_impl(ctx):
  98     trans_srcs = _get_transitive_srcs(ctx.files.srcs, ctx.attr.deps)
  99     trans_includes = _get_transitive_includes(
 100         _resolve_includes(ctx, ctx.attr.includes),
 101         ctx.attr.deps,
 102     )
 103
 104     # Note that we include srcs in runfiles. A td_library doesn't compile to
 105     # produce an output: it's just a depset of source files and include
 106     # directories. So if it is needed for execution of some rule (likely
 107     # something running tblgen as a test action), the files needed are the same
 108     # as the source files.
 109     # Note: not using merge_all, as that is not available in Bazel 4.0
 110     runfiles = ctx.runfiles(ctx.files.srcs)
 111     for src in ctx.attr.srcs:
 112         runfiles = runfiles.merge(src[DefaultInfo].default_runfiles)
 113     for dep in ctx.attr.deps:
 114         runfiles = runfiles.merge(dep[DefaultInfo].default_runfiles)
 115
 116     return [
 117         DefaultInfo(files = trans_srcs, runfiles = runfiles),
 118         TdInfo(
 119             transitive_sources = trans_srcs,
 120             transitive_includes = trans_includes,
 121         ),
 122     ]
 123
 124 td_library = rule(
 125     _td_library_impl,
 126     attrs = {
 127         "srcs": attr.label_list(allow_files = True),
 128         "includes": attr.string_list(
 129             doc = "Include paths to be added to the final TableGen tool" +
 130                   " invocation. Relative paths are interpreted as relative to" +
 131                   " the current label's package. Absolute paths are" +
 132                   " interpreted as relative to the current label's workspace",
 133         ),
 134         # TODO(gcmn): limit to TdInfo providers.
 135         "deps": attr.label_list(
 136             doc = "Dependencies providing TableGen source files and include" +
 137                   " paths.",
 138         ),
 139     },
 140 )
 141
 142 def _gentbl_rule_impl(ctx):
 143     td_file = ctx.file.td_file
 144
 145     trans_srcs = _get_transitive_srcs(
 146         ctx.files.td_srcs + [td_file],
 147         ctx.attr.deps,
 148     )
 149
 150     # Note that the td_file.dirname is already relative to the execution root,
 151     # i.e. may contain an `external/<workspace_name>` prefix if the current
 152     # workspace is not the main workspace. Therefore it is not included in the
 153     # _resolve_includes call that prepends this prefix.
 154     trans_includes = _get_transitive_includes(
 155         _resolve_includes(ctx, ctx.attr.includes + ["/"]) +
 156         _prefix_roots(ctx, [td_file.dirname]),
 157         ctx.attr.deps,
 158     )
 159
 160     args = ctx.actions.args()
 161     args.add_all(ctx.attr.opts)
 162     args.add(td_file)
 163     args.add_all(trans_includes, before_each = "-I")
 164
 165     args.add("-o", ctx.outputs.out.path)
 166
 167     ctx.actions.run(
 168         outputs = [ctx.outputs.out],
 169         inputs = trans_srcs,
 170         executable = ctx.executable.tblgen,
 171         arguments = [args],
 172         # Make sure action_env settings are honored so the env is the same as
 173         # when the tool was built. Important for locating shared libraries with
 174         # a custom LD_LIBRARY_PATH.
 175         use_default_shell_env = True,
 176         mnemonic = "TdGenerate",
 177     )
 178
 179     return [DefaultInfo()]
 180
 181 gentbl_rule = rule(
 182     _gentbl_rule_impl,
 183     doc = "Generates tabular code from a table definition file.",
 184     attrs = {
 185         "tblgen": attr.label(
 186             doc = "The TableGen executable with which to generate `out`.",
 187             executable = True,
 188             cfg = "exec",
 189         ),
 190         "td_file": attr.label(
 191             doc = "The TableGen file to run through `tblgen`.",
 192             allow_single_file = True,
 193             mandatory = True,
 194         ),
 195         "td_srcs": attr.label_list(
 196             doc = "Additional TableGen files included by `td_file`. It is not" +
 197                   " necessary to list td_file here (though not an error).",
 198             allow_files = True,
 199         ),
 200         # TODO(gcmn): limit to TdInfo providers.
 201         "deps": attr.label_list(
 202             doc = "Dependencies providing TableGen source files and include" +
 203                   " paths.",
 204         ),
 205         "out": attr.output(
 206             doc = "The output file for the TableGen invocation.",
 207             mandatory = True,
 208         ),
 209         "opts": attr.string_list(
 210             doc = "Additional command line options to add to the TableGen" +
 211                   " invocation. For include arguments, prefer to use" +
 212                   " `includes`.",
 213         ),
 214         "includes": attr.string_list(
 215             doc = "Include paths to be added to the final TableGen tool" +
 216                   " invocation. Relative paths are interpreted as relative to" +
 217                   " the current label's package. Absolute paths are" +
 218                   " interpreted as relative to the current label's workspace." +
 219                   " Includes are applied from all roots available in the" +
 220                   " execution environment (source, genfiles, and bin" +
 221                   " directories). The execution roots themselves and the " +
 222                   " directory of td_file are always added.",
 223         ),
 224     },
 225 )
 226
 227 # TODO(gcmn): Figure out how to reduce duplication with _gentbl_rule_impl
 228 def _gentbl_test_impl(ctx):
 229     td_file = ctx.file.td_file
 230
 231     # Note that the td_file.dirname is already relative to the execution root,
 232     # i.e. may contain an `external/<workspace_name>` prefix if the current
 233     # workspace is not the main workspace. Therefore it is not included in the
 234     # _resolve_includes call that prepends this prefix.
 235     trans_includes = _get_transitive_includes(
 236         _resolve_includes(ctx, ctx.attr.includes + ["/"]) +
 237         _prefix_roots(ctx, [td_file.dirname]),
 238         ctx.attr.deps,
 239     )
 240
 241     test_args = [ctx.executable.tblgen.short_path]
 242     test_args.extend(ctx.attr.opts)
 243     test_args.append(td_file.path)
 244     test_args.extend(["-I " + include for include in trans_includes.to_list()])
 245
 246     test_args.extend(["-o", "/dev/null"])
 247
 248     ctx.actions.write(
 249         ctx.outputs.executable,
 250         content = " ".join(test_args),
 251         is_executable = True,
 252     )
 253
 254     # Note: not using merge_all, as that is not available in Bazel 4.0
 255     runfiles = ctx.runfiles(
 256         files = [ctx.executable.tblgen],
 257         transitive_files = _get_transitive_srcs(
 258             ctx.files.td_srcs + [td_file],
 259             ctx.attr.deps,
 260         ),
 261     )
 262     for src in ctx.attr.td_srcs:
 263         runfiles = runfiles.merge(src[DefaultInfo].default_runfiles)
 264     for dep in ctx.attr.deps:
 265         runfiles = runfiles.merge(dep[DefaultInfo].default_runfiles)
 266
 267     return [
 268         coverage_common.instrumented_files_info(
 269             ctx,
 270             source_attributes = ["td_file", "td_srcs"],
 271             dependency_attributes = ["tblgen", "deps"],
 272         ),
 273         DefaultInfo(runfiles = runfiles),
 274     ]
 275
 276 gentbl_test = rule(
 277     _gentbl_test_impl,
 278     test = True,
 279     doc = "A shell test that tests the given TablegGen invocation. Note" +
 280           " that unlike gentbl_rule, this builds and invokes `tblgen` in the" +
 281           " target configuration. Takes all the same arguments as gentbl_rule" +
 282           " except for `out` (as it does not generate any output)",
 283     attrs = {
 284         "tblgen": attr.label(
 285             doc = "The TableGen executable run in the shell command. Note" +
 286                   " that this is built in the target configuration.",
 287             executable = True,
 288             cfg = "target",
 289         ),
 290         "td_file": attr.label(
 291             doc = "See gentbl_rule.td_file",
 292             allow_single_file = True,
 293             mandatory = True,
 294         ),
 295         "td_srcs": attr.label_list(
 296             doc = "See gentbl_rule.td_srcs",
 297             allow_files = True,
 298         ),
 299         "deps": attr.label_list(doc = "See gentbl_rule.deps"),
 300         "opts": attr.string_list(doc = "See gentbl_rule.opts"),
 301         "includes": attr.string_list(doc = "See gentbl_rule.includes"),
 302     },
 303 )
 304
 305 def gentbl_filegroup(
 306         name,
 307         tblgen,
 308         td_file,
 309         tbl_outs,
 310         td_srcs = [],
 311         includes = [],
 312         deps = [],
 313         test = False,
 314         skip_opts = [],
 315         **kwargs):
 316     """Create multiple TableGen generated files using the same tool and input.
 317
 318     All generated outputs are bundled in a file group with the given name.
 319
 320     Args:
 321       name: The name of the generated filegroup rule for use in dependencies.
 322       tblgen: The binary used to produce the output.
 323       td_file: The primary table definitions file.
 324       tbl_outs: A list of tuples ([opts], out), where each 'opts' is a list of
 325         options passed to tblgen, each option being a string, and 'out' is the
 326         corresponding output file produced.
 327       td_srcs: See gentbl_rule.td_srcs
 328       includes: See gentbl_rule.includes
 329       deps: See gentbl_rule.deps
 330       test: Whether to create a shell test that invokes the tool too.
 331       skip_opts: Files generated using these opts in tbl_outs will be excluded
 332         from the generated filegroup.
 333       **kwargs: Extra keyword arguments to pass to all generated rules.
 334     """
 335
 336     for (opts, out) in tbl_outs:
 337         first_opt = opts[0] if opts else ""
 338         rule_suffix = "_{}_{}".format(
 339             first_opt.replace("-", "_").replace("=", "_"),
 340             str(hash(" ".join(opts))),
 341         )
 342         gentbl_name = "%s_%s_genrule" % (name, rule_suffix)
 343         gentbl_rule(
 344             name = gentbl_name,
 345             td_file = td_file,
 346             tblgen = tblgen,
 347             opts = opts,
 348             td_srcs = td_srcs,
 349             deps = deps,
 350             includes = includes,
 351             out = out,
 352             **kwargs
 353         )
 354
 355         if test:
 356             # Also run the generator in the target configuration as a test. This
 357             # means it gets run with asserts and sanitizers and such when they
 358             # are enabled and is counted in coverage.
 359             gentbl_test(
 360                 name = "%s_test" % (gentbl_name,),
 361                 td_file = td_file,
 362                 tblgen = tblgen,
 363                 opts = opts,
 364                 td_srcs = td_srcs,
 365                 deps = deps,
 366                 includes = includes,
 367                 # Shell files not executable on Windows.
 368                 # TODO(gcmn): Support windows.
 369                 tags = ["no_windows"],
 370                 **kwargs
 371             )
 372
 373     included_srcs = [f for (opts, f) in tbl_outs if not any([skip_opt in opts for skip_opt in skip_opts])]
 374     native.filegroup(
 375         name = name,
 376         srcs = included_srcs,
 377         **kwargs
 378     )
 379
 380 def gentbl_cc_library(
 381         name,
 382         tblgen,
 383         td_file,
 384         tbl_outs,
 385         td_srcs = [],
 386         includes = [],
 387         deps = [],
 388         strip_include_prefix = None,
 389         test = False,
 390         copts = None,
 391         **kwargs):
 392     """Create multiple TableGen generated files using the same tool and input.
 393
 394     All generated outputs are bundled in a cc_library rule.
 395
 396     Args:
 397       name: The name of the generated cc_library rule for use in dependencies.
 398       tblgen: The binary used to produce the output.
 399       td_file: The primary table definitions file.
 400       tbl_outs: A list of tuples ([opts], out), where each 'opts' is a list of
 401         options passed to tblgen, each option being a string, and 'out' is the
 402         corresponding output file produced.
 403       td_srcs: See gentbl_rule.td_srcs
 404       includes: See gentbl_rule.includes
 405       deps: See gentbl_rule.deps
 406       strip_include_prefix: attribute to pass through to cc_library.
 407       test: whether to create a shell test that invokes the tool too.
 408       copts: list of copts to pass to cc_library.
 409       **kwargs: Extra keyword arguments to pass to all generated rules.
 410     """
 411
 412     filegroup_name = name + "_filegroup"
 413     gentbl_filegroup(
 414         name = filegroup_name,
 415         tblgen = tblgen,
 416         td_file = td_file,
 417         tbl_outs = tbl_outs,
 418         td_srcs = td_srcs,
 419         includes = includes,
 420         deps = deps,
 421         test = test,
 422         skip_opts = ["-gen-op-doc"],
 423         **kwargs
 424     )
 425     native.cc_library(
 426         name = name,
 427         # strip_include_prefix does not apply to textual_hdrs.
 428         # https://github.com/bazelbuild/bazel/issues/12424
 429         hdrs = [":" + filegroup_name] if strip_include_prefix else [],
 430         strip_include_prefix = strip_include_prefix,
 431         textual_hdrs = [":" + filegroup_name],
 432         copts = copts,
 433         **kwargs
 434     )
 435
 436 def _gentbl_shard_impl(ctx):
 437     args = ctx.actions.args()
 438     args.add(ctx.file.src_file)
 439     args.add("-op-shard-index", ctx.attr.index)
 440     args.add("-o", ctx.outputs.out.path)
 441     ctx.actions.run(
 442         outputs = [ctx.outputs.out],
 443         inputs = [ctx.file.src_file],
 444         executable = ctx.executable.sharder,
 445         arguments = [args],
 446         use_default_shell_env = True,
 447         mnemonic = "ShardGenerate",
 448     )
 449
 450 gentbl_shard_rule = rule(
 451     _gentbl_shard_impl,
 452     doc = "",
 453     output_to_genfiles = True,
 454     attrs = {
 455         "index": attr.int(mandatory = True, doc = ""),
 456         "sharder": attr.label(
 457             doc = "",
 458             executable = True,
 459             cfg = "exec",
 460         ),
 461         "src_file": attr.label(
 462             doc = "",
 463             allow_single_file = True,
 464             mandatory = True,
 465         ),
 466         "out": attr.output(
 467             doc = "",
 468             mandatory = True,
 469         ),
 470     },
 471 )
 472
 473 def gentbl_sharded_ops(
 474         name,
 475         tblgen,
 476         sharder,
 477         td_file,
 478         shard_count,
 479         src_file,
 480         src_out,
 481         hdr_out,
 482         test = False,
 483         includes = [],
 484         strip_include_prefix = None,
 485         deps = []):
 486     """Generate sharded op declarations and definitions.
 487
 488     This special build rule shards op definitions in a TableGen file and generates multiple copies
 489     of a template source file for including and compiling each shard. The rule defines a filegroup
 490     consisting of the source shards, the generated source file, and the generated header file.
 491
 492     Args:
 493       name: The name of the filegroup.
 494       tblgen: The binary used to produce the output.
 495       sharder: The source file sharder to use.
 496       td_file: The primary table definitions file.
 497       shard_count: The number of op definition shards to produce.
 498       src_file: The source file template.
 499       src_out: The generated source file.
 500       hdr_out: The generated header file.
 501       test: Whether this is a test target.
 502       includes: See gentbl_rule.includes
 503       deps: See gentbl_rule.deps
 504       strip_include_prefix: Attribute to pass through to cc_library.
 505     """
 506     cc_lib_name = name + "__gentbl_cc_lib"
 507     gentbl_cc_library(
 508         name = cc_lib_name,
 509         strip_include_prefix = strip_include_prefix,
 510         includes = includes,
 511         tbl_outs = [
 512             (
 513                 [
 514                     "-gen-op-defs",
 515                     "-op-shard-count=" + str(shard_count),
 516                 ],
 517                 src_out,
 518             ),
 519             (
 520                 [
 521                     "-gen-op-decls",
 522                     "-op-shard-count=" + str(shard_count),
 523                 ],
 524                 hdr_out,
 525             ),
 526         ],
 527         tblgen = tblgen,
 528         td_file = td_file,
 529         test = test,
 530         deps = deps,
 531     )
 532     all_files = [hdr_out, src_out]
 533     for i in range(0, shard_count):
 534         out_file = "shard_copy_" + str(i) + "_" + src_file
 535         gentbl_shard_rule(
 536             index = i,
 537             name = name + "__src_shard" + str(i),
 538             testonly = test,
 539             out = out_file,
 540             sharder = sharder,
 541             src_file = src_file,
 542         )
 543         all_files.append(out_file)
 544     native.filegroup(name = name, srcs = all_files)
 545
 546 def gentbl_sharded_op_defs(name, source_file, shard_count):
 547     """Generates multiple copies of a source file that includes sharded op definitions.
 548
 549     Args:
 550       name: The name of the rule.
 551       source_file: The source to copy.
 552       shard_count: The number of shards.
 553
 554     Returns:
 555       A list of the copied filenames to be included in the dialect library.
 556     """
 557     copies = []
 558     for i in range(0, shard_count):
 559         out_file = "shard_copy_" + str(i) + "_" + source_file
 560         copies.append(out_file)
 561         native.genrule(
 562             name = name + "_shard_" + str(i),
 563             srcs = [source_file],
 564             outs = [out_file],
 565             cmd = "echo -e \"#define GET_OP_DEFS_" + str(i) + "\n$$(cat $(SRCS))\" > $(OUTS)",
 566         )
 567     return copies