From 546df5f3729c26985405d9f93d8a36bd1a54778b Mon Sep 17 00:00:00 2001 From: Keith Smiley Date: Fri, 12 Sep 2025 19:42:48 +0000 Subject: [PATCH 1/2] [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. --- .../llvm-project-overlay/mlir/tblgen.bzl | 38 +++++++------------ 1 file changed, 14 insertions(+), 24 deletions(-) diff --git a/utils/bazel/llvm-project-overlay/mlir/tblgen.bzl b/utils/bazel/llvm-project-overlay/mlir/tblgen.bzl index c94935216e0e9..c5c4d544691b6 100644 --- a/utils/bazel/llvm-project-overlay/mlir/tblgen.bzl +++ b/utils/bazel/llvm-project-overlay/mlir/tblgen.bzl @@ -62,20 +62,6 @@ def _get_transitive_includes(includes, deps): transitive = [_get_dep_transitive_includes(dep) for dep in deps], ) -def _prefix_roots(ctx, includes): - """Map the given includes to be relative to all root directories. - - This will expand them to be relative to all the root directories available - in the execution environment for ctx.run (bin and genfiles in addition to - the normal source root) - """ - prefixed_includes = [] - for include in includes: - prefixed_includes.append(include) - prefixed_includes.append(paths.join(ctx.genfiles_dir.path, include)) - prefixed_includes.append(paths.join(ctx.bin_dir.path, include)) - return prefixed_includes - def _resolve_includes(ctx, includes): """Resolves include paths to paths relative to the execution root. @@ -92,7 +78,7 @@ def _resolve_includes(ctx, includes): else: include = paths.join(package, include) include = paths.join(workspace_root, include) - resolved_includes.extend(_prefix_roots(ctx, [include])) + resolved_includes.append(include) return resolved_includes def _td_library_impl(ctx): @@ -140,6 +126,9 @@ td_library = rule( }, ) +def _format_includes(output): + return lambda x: ["-I", x, "-I", paths.join(output.root.path, x)] + def _gentbl_rule_impl(ctx): td_file = ctx.file.td_file @@ -153,22 +142,21 @@ def _gentbl_rule_impl(ctx): # workspace is not the main workspace. Therefore it is not included in the # _resolve_includes call that prepends this prefix. trans_includes = _get_transitive_includes( - _resolve_includes(ctx, ctx.attr.includes + ["/"]) + - _prefix_roots(ctx, [td_file.dirname]), + _resolve_includes(ctx, ctx.attr.includes + ["/"]) + [td_file.dirname], ctx.attr.deps, ) args = ctx.actions.args() args.add_all(ctx.attr.opts) args.add(td_file) - args.add_all(trans_includes, before_each = "-I") - - args.add("-o", ctx.outputs.out.path) + args.add_all(trans_includes, map_each = _format_includes(ctx.outputs.out), allow_closure = True) + args.add("-o", ctx.outputs.out) ctx.actions.run( outputs = [ctx.outputs.out], inputs = trans_srcs, executable = ctx.executable.tblgen, + execution_requirements = {"supports-path-mapping": "1"}, arguments = [args], # Make sure action_env settings are honored so the env is the same as # when the tool was built. Important for locating shared libraries with @@ -234,15 +222,16 @@ def _gentbl_test_impl(ctx): # workspace is not the main workspace. Therefore it is not included in the # _resolve_includes call that prepends this prefix. trans_includes = _get_transitive_includes( - _resolve_includes(ctx, ctx.attr.includes + ["/"]) + - _prefix_roots(ctx, [td_file.dirname]), + _resolve_includes(ctx, ctx.attr.includes + ["/"]) + [td_file.dirname], ctx.attr.deps, ) test_args = [ctx.executable.tblgen.short_path] test_args.extend(ctx.attr.opts) test_args.append(td_file.path) - test_args.extend(["-I " + include for include in trans_includes.to_list()]) + for include in trans_includes.to_list(): + test_args.extend(["-I", include]) + test_args.extend(["-I", paths.join(ctx.bin_dir.path, include)]) test_args.extend(["-o", "/dev/null"]) @@ -440,11 +429,12 @@ def _gentbl_shard_impl(ctx): args = ctx.actions.args() args.add(ctx.file.src_file) args.add("-op-shard-index", ctx.attr.index) - args.add("-o", ctx.outputs.out.path) + args.add("-o", ctx.outputs.out) ctx.actions.run( outputs = [ctx.outputs.out], inputs = [ctx.file.src_file], executable = ctx.executable.sharder, + execution_requirements = {"supports-path-mapping": "1"}, arguments = [args], use_default_shell_env = True, mnemonic = "ShardGenerate", From 3460ab7aec7a51d8896d6b5704c77784b347591f Mon Sep 17 00:00:00 2001 From: Keith Smiley Date: Mon, 15 Sep 2025 16:13:13 +0000 Subject: [PATCH 2/2] reduce extend calls --- utils/bazel/llvm-project-overlay/mlir/tblgen.bzl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/utils/bazel/llvm-project-overlay/mlir/tblgen.bzl b/utils/bazel/llvm-project-overlay/mlir/tblgen.bzl index c5c4d544691b6..35888aac37e17 100644 --- a/utils/bazel/llvm-project-overlay/mlir/tblgen.bzl +++ b/utils/bazel/llvm-project-overlay/mlir/tblgen.bzl @@ -229,9 +229,11 @@ def _gentbl_test_impl(ctx): test_args = [ctx.executable.tblgen.short_path] test_args.extend(ctx.attr.opts) test_args.append(td_file.path) - for include in trans_includes.to_list(): - test_args.extend(["-I", include]) - test_args.extend(["-I", paths.join(ctx.bin_dir.path, include)]) + test_args.extend([ + arg + for include in trans_includes.to_list() + for arg in ["-I", include, "-I", paths.join(ctx.bin_dir.path, include)] + ]) test_args.extend(["-o", "/dev/null"])