Skip to content

Commit 546df5f

Browse files
committed
[bazel] Update tblgen rules to support path-mapping
Path mapping is a feature of bazel that allows actions to be deduplicated across bazel transitions if they are otherwise identical. This is helpful if you build a binary in N transitions, where the settings differences are irrelevant to this action. In our case we build multiple python native extensions transitioning on the python version they target, and before this change would run each of these actions once per python version even though the outputs would be identical. This is a no-op unless `--experimental_output_paths=strip` is passed. The changes here are just enough to make bazel automatically remap the paths, which is done by how you use the args object. The core change is that instead of carrying around paths that have `ctx.bin_dir` hardcoded in the strings. This is done by mapping them with the output file object's root path when adding them to the args. As a side effect this drops the genfiles_dir, but that has been the same as bin_dir for a long time so hopefully that's a no-op for folks.
1 parent 1756b6e commit 546df5f

File tree

1 file changed

+14
-24
lines changed

1 file changed

+14
-24
lines changed

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

Lines changed: 14 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -62,20 +62,6 @@ def _get_transitive_includes(includes, deps):
6262
transitive = [_get_dep_transitive_includes(dep) for dep in deps],
6363
)
6464

65-
def _prefix_roots(ctx, includes):
66-
"""Map the given includes to be relative to all root directories.
67-
68-
This will expand them to be relative to all the root directories available
69-
in the execution environment for ctx.run (bin and genfiles in addition to
70-
the normal source root)
71-
"""
72-
prefixed_includes = []
73-
for include in includes:
74-
prefixed_includes.append(include)
75-
prefixed_includes.append(paths.join(ctx.genfiles_dir.path, include))
76-
prefixed_includes.append(paths.join(ctx.bin_dir.path, include))
77-
return prefixed_includes
78-
7965
def _resolve_includes(ctx, includes):
8066
"""Resolves include paths to paths relative to the execution root.
8167
@@ -92,7 +78,7 @@ def _resolve_includes(ctx, includes):
9278
else:
9379
include = paths.join(package, include)
9480
include = paths.join(workspace_root, include)
95-
resolved_includes.extend(_prefix_roots(ctx, [include]))
81+
resolved_includes.append(include)
9682
return resolved_includes
9783

9884
def _td_library_impl(ctx):
@@ -140,6 +126,9 @@ td_library = rule(
140126
},
141127
)
142128

129+
def _format_includes(output):
130+
return lambda x: ["-I", x, "-I", paths.join(output.root.path, x)]
131+
143132
def _gentbl_rule_impl(ctx):
144133
td_file = ctx.file.td_file
145134

@@ -153,22 +142,21 @@ def _gentbl_rule_impl(ctx):
153142
# workspace is not the main workspace. Therefore it is not included in the
154143
# _resolve_includes call that prepends this prefix.
155144
trans_includes = _get_transitive_includes(
156-
_resolve_includes(ctx, ctx.attr.includes + ["/"]) +
157-
_prefix_roots(ctx, [td_file.dirname]),
145+
_resolve_includes(ctx, ctx.attr.includes + ["/"]) + [td_file.dirname],
158146
ctx.attr.deps,
159147
)
160148

161149
args = ctx.actions.args()
162150
args.add_all(ctx.attr.opts)
163151
args.add(td_file)
164-
args.add_all(trans_includes, before_each = "-I")
165-
166-
args.add("-o", ctx.outputs.out.path)
152+
args.add_all(trans_includes, map_each = _format_includes(ctx.outputs.out), allow_closure = True)
153+
args.add("-o", ctx.outputs.out)
167154

168155
ctx.actions.run(
169156
outputs = [ctx.outputs.out],
170157
inputs = trans_srcs,
171158
executable = ctx.executable.tblgen,
159+
execution_requirements = {"supports-path-mapping": "1"},
172160
arguments = [args],
173161
# Make sure action_env settings are honored so the env is the same as
174162
# when the tool was built. Important for locating shared libraries with
@@ -234,15 +222,16 @@ def _gentbl_test_impl(ctx):
234222
# workspace is not the main workspace. Therefore it is not included in the
235223
# _resolve_includes call that prepends this prefix.
236224
trans_includes = _get_transitive_includes(
237-
_resolve_includes(ctx, ctx.attr.includes + ["/"]) +
238-
_prefix_roots(ctx, [td_file.dirname]),
225+
_resolve_includes(ctx, ctx.attr.includes + ["/"]) + [td_file.dirname],
239226
ctx.attr.deps,
240227
)
241228

242229
test_args = [ctx.executable.tblgen.short_path]
243230
test_args.extend(ctx.attr.opts)
244231
test_args.append(td_file.path)
245-
test_args.extend(["-I " + include for include in trans_includes.to_list()])
232+
for include in trans_includes.to_list():
233+
test_args.extend(["-I", include])
234+
test_args.extend(["-I", paths.join(ctx.bin_dir.path, include)])
246235

247236
test_args.extend(["-o", "/dev/null"])
248237

@@ -440,11 +429,12 @@ def _gentbl_shard_impl(ctx):
440429
args = ctx.actions.args()
441430
args.add(ctx.file.src_file)
442431
args.add("-op-shard-index", ctx.attr.index)
443-
args.add("-o", ctx.outputs.out.path)
432+
args.add("-o", ctx.outputs.out)
444433
ctx.actions.run(
445434
outputs = [ctx.outputs.out],
446435
inputs = [ctx.file.src_file],
447436
executable = ctx.executable.sharder,
437+
execution_requirements = {"supports-path-mapping": "1"},
448438
arguments = [args],
449439
use_default_shell_env = True,
450440
mnemonic = "ShardGenerate",

0 commit comments

Comments
 (0)