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