Skip to content

Commit 84eb453

Browse files
draganmladjenovicdraganmladjenovic
committed
Add hlo runner with TE support
bazel build --config=rocm_clang_official --repo_env=HERMETIC_PYTHON_VERSION=$(python --version | grep -o "3.[0-9]*") --disk_cache=/tf/cache //xla/tools/multihost_hlo_runner:python_hlo_runner_main && PYTHONPATH=$(bazel info bazel-bin 2>/dev/null)/xla/tools/multihost_hlo_runner/python_hlo_runner_main.runfiles/xla python xla/tools/multihost_hlo_runner/python_hlo_runner_main.py atention.hlo
1 parent ff16803 commit 84eb453

File tree

9 files changed

+934
-14
lines changed

9 files changed

+934
-14
lines changed

‎WORKSPACE‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,9 @@ load("//third_party/py:python_init_repositories.bzl", "python_init_repositories"
1717

1818
python_init_repositories(
1919
requirements = {
20+
"3.10": "//:requirements_lock_3_10.txt",
2021
"3.11": "//:requirements_lock_3_11.txt",
22+
"3.12": "//:requirements_lock_3_12.txt",
2123
},
2224
)
2325

‎requirements_lock_3_10.txt‎

Lines changed: 707 additions & 0 deletions
Large diffs are not rendered by default.

‎requirements_lock_3_12.txt‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
numpy==2.1.3 \
2+
--hash=sha256:016d0f6f5e77b0f0d45d77387ffa4bb89816b57c835580c3ce8e099ef830befe \
3+
--hash=sha256:02135ade8b8a84011cbb67dc44e07c58f28575cf9ecf8ab304e51c05528c19f0 \
4+
--hash=sha256:08788d27a5fd867a663f6fc753fd7c3ad7e92747efc73c53bca2f19f8bc06f48 \
5+
--hash=sha256:0d30c543f02e84e92c4b1f415b7c6b5326cbe45ee7882b6b77db7195fb971e3a \
6+
--hash=sha256:0fa14563cc46422e99daef53d725d0c326e99e468a9320a240affffe87852564 \
7+
--hash=sha256:13138eadd4f4da03074851a698ffa7e405f41a0845a6b1ad135b81596e4e9958 \
8+
--hash=sha256:14e253bd43fc6b37af4921b10f6add6925878a42a0c5fe83daee390bca80bc17 \
9+
--hash=sha256:15cb89f39fa6d0bdfb600ea24b250e5f1a3df23f901f51c8debaa6a5d122b2f0 \
10+
--hash=sha256:17ee83a1f4fef3c94d16dc1802b998668b5419362c8a4f4e8a491de1b41cc3ee \
11+
--hash=sha256:2312b2aa89e1f43ecea6da6ea9a810d06aae08321609d8dc0d0eda6d946a541b \
12+
--hash=sha256:2564fbdf2b99b3f815f2107c1bbc93e2de8ee655a69c261363a1172a79a257d4 \
13+
--hash=sha256:3522b0dfe983a575e6a9ab3a4a4dfe156c3e428468ff08ce582b9bb6bd1d71d4 \
14+
--hash=sha256:4394bc0dbd074b7f9b52024832d16e019decebf86caf909d94f6b3f77a8ee3b6 \
15+
--hash=sha256:45966d859916ad02b779706bb43b954281db43e185015df6eb3323120188f9e4 \
16+
--hash=sha256:4d1167c53b93f1f5d8a139a742b3c6f4d429b54e74e6b57d0eff40045187b15d \
17+
--hash=sha256:4f2015dfe437dfebbfce7c85c7b53d81ba49e71ba7eadbf1df40c915af75979f \
18+
--hash=sha256:50ca6aba6e163363f132b5c101ba078b8cbd3fa92c7865fd7d4d62d9779ac29f \
19+
--hash=sha256:50d18c4358a0a8a53f12a8ba9d772ab2d460321e6a93d6064fc22443d189853f \
20+
--hash=sha256:5641516794ca9e5f8a4d17bb45446998c6554704d888f86df9b200e66bdcce56 \
21+
--hash=sha256:576a1c1d25e9e02ed7fa5477f30a127fe56debd53b8d2c89d5578f9857d03ca9 \
22+
--hash=sha256:6a4825252fcc430a182ac4dee5a505053d262c807f8a924603d411f6718b88fd \
23+
--hash=sha256:72dcc4a35a8515d83e76b58fdf8113a5c969ccd505c8a946759b24e3182d1f23 \
24+
--hash=sha256:747641635d3d44bcb380d950679462fae44f54b131be347d5ec2bce47d3df9ed \
25+
--hash=sha256:762479be47a4863e261a840e8e01608d124ee1361e48b96916f38b119cfda04a \
26+
--hash=sha256:78574ac2d1a4a02421f25da9559850d59457bac82f2b8d7a44fe83a64f770098 \
27+
--hash=sha256:825656d0743699c529c5943554d223c021ff0494ff1442152ce887ef4f7561a1 \
28+
--hash=sha256:8637dcd2caa676e475503d1f8fdb327bc495554e10838019651b76d17b98e512 \
29+
--hash=sha256:96fe52fcdb9345b7cd82ecd34547fca4321f7656d500eca497eb7ea5a926692f \
30+
--hash=sha256:973faafebaae4c0aaa1a1ca1ce02434554d67e628b8d805e61f874b84e136b09 \
31+
--hash=sha256:996bb9399059c5b82f76b53ff8bb686069c05acc94656bb259b1d63d04a9506f \
32+
--hash=sha256:a38c19106902bb19351b83802531fea19dee18e5b37b36454f27f11ff956f7fc \
33+
--hash=sha256:a6b46587b14b888e95e4a24d7b13ae91fa22386c199ee7b418f449032b2fa3b8 \
34+
--hash=sha256:a9f7f672a3388133335589cfca93ed468509cb7b93ba3105fce780d04a6576a0 \
35+
--hash=sha256:aa08e04e08aaf974d4458def539dece0d28146d866a39da5639596f4921fd761 \
36+
--hash=sha256:b0df3635b9c8ef48bd3be5f862cf71b0a4716fa0e702155c45067c6b711ddcef \
37+
--hash=sha256:b47fbb433d3260adcd51eb54f92a2ffbc90a4595f8970ee00e064c644ac788f5 \
38+
--hash=sha256:baed7e8d7481bfe0874b566850cb0b85243e982388b7b23348c6db2ee2b2ae8e \
39+
--hash=sha256:bc6f24b3d1ecc1eebfbf5d6051faa49af40b03be1aaa781ebdadcbc090b4539b \
40+
--hash=sha256:c006b607a865b07cd981ccb218a04fc86b600411d83d6fc261357f1c0966755d \
41+
--hash=sha256:c181ba05ce8299c7aa3125c27b9c2167bca4a4445b7ce73d5febc411ca692e43 \
42+
--hash=sha256:c7662f0e3673fe4e832fe07b65c50342ea27d989f92c80355658c7f888fcc83c \
43+
--hash=sha256:c80e4a09b3d95b4e1cac08643f1152fa71a0a821a2d4277334c88d54b2219a41 \
44+
--hash=sha256:c894b4305373b9c5576d7a12b473702afdf48ce5369c074ba304cc5ad8730dff \
45+
--hash=sha256:d7aac50327da5d208db2eec22eb11e491e3fe13d22653dce51b0f4109101b408 \
46+
--hash=sha256:d89dd2b6da69c4fff5e39c28a382199ddedc3a5be5390115608345dec660b9e2 \
47+
--hash=sha256:d9beb777a78c331580705326d2367488d5bc473b49a9bc3036c154832520aca9 \
48+
--hash=sha256:dc258a761a16daa791081d026f0ed4399b582712e6fc887a95af09df10c5ca57 \
49+
--hash=sha256:e14e26956e6f1696070788252dcdff11b4aca4c3e8bd166e0df1bb8f315a67cb \
50+
--hash=sha256:e6988e90fcf617da2b5c78902fe8e668361b43b4fe26dbf2d7b0f8034d4cafb9 \
51+
--hash=sha256:e711e02f49e176a01d0349d82cb5f05ba4db7d5e7e0defd026328e5cfb3226d3 \
52+
--hash=sha256:ea4dedd6e394a9c180b33c2c872b92f7ce0f8e7ad93e9585312b0c5a04777a4a \
53+
--hash=sha256:ecc76a9ba2911d8d37ac01de72834d8849e55473457558e12995f4cd53e778e0 \
54+
--hash=sha256:f55ba01150f52b1027829b50d70ef1dafd9821ea82905b63936668403c3b471e \
55+
--hash=sha256:f653490b33e9c3a4c1c01d41bc2aef08f9475af51146e4a7710c450cf9761598 \
56+
--hash=sha256:fa2d1337dc61c8dc417fbccf20f6d1e139896a30721b7f1e832b2bb6ef4eb6c4
57+
lit==17.0.6 \
58+
--hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b
59+
ml-dtypes==0.5.3 \
60+
--hash=sha256:01de48de4537dc3c46e684b969a40ec36594e7eeb7c69e9a093e7239f030a28a \
61+
--hash=sha256:0a1d68a7cb53e3f640b2b6a34d12c0542da3dd935e560fdf463c0c77f339fc20 \
62+
--hash=sha256:0cd5a6c711b5350f3cbc2ac28def81cd1c580075ccb7955e61e9d8f4bfd40d24 \
63+
--hash=sha256:0e44a3761f64bc009d71ddb6d6c71008ba21b53ab6ee588dadab65e2fa79eafc \
64+
--hash=sha256:156418abeeda48ea4797db6776db3c5bdab9ac7be197c1233771e0880c304057 \
65+
--hash=sha256:19f6c3a4f635c2fc9e2aa7d91416bd7a3d649b48350c51f7f715a09370a90d93 \
66+
--hash=sha256:1b255acada256d1fa8c35ed07b5f6d18bc21d1556f842fbc2d5718aea2cd9e55 \
67+
--hash=sha256:1db60c154989af253f6c4a34e8a540c2c9dce4d770784d426945e09908fbb177 \
68+
--hash=sha256:2db74788fc01914a3c7f7da0763427280adfc9cd377e9604b6b64eb8097284bd \
69+
--hash=sha256:4a177b882667c69422402df6ed5c3428ce07ac2c1f844d8a1314944651439458 \
70+
--hash=sha256:4cae435a68861660af81fa3c5af16b70ca11a17275c5b662d9c6f58294e0f113 \
71+
--hash=sha256:5103856a225465371fe119f2fef737402b705b810bd95ad5f348e6e1a6ae21af \
72+
--hash=sha256:58e39349d820b5702bb6f94ea0cb2dc8ec62ee81c0267d9622067d8333596a46 \
73+
--hash=sha256:5ab039ffb40f3dc0aeeeba84fd6c3452781b5e15bef72e2d10bcb33e4bbffc39 \
74+
--hash=sha256:5ee72568d46b9533ad54f78b1e1f3067c0534c5065120ea8ecc6f210d22748b3 \
75+
--hash=sha256:66c2756ae6cfd7f5224e355c893cfd617fa2f747b8bbd8996152cbdebad9a184 \
76+
--hash=sha256:6936283b56d74fbec431ca57ce58a90a908fdbd14d4e2d22eea6d72bb208a7b7 \
77+
--hash=sha256:8b1a6e231b0770f2894910f1dce6d2f31d65884dbf7668f9b08d73623cdca909 \
78+
--hash=sha256:8bb9cd1ce63096567f5f42851f5843b5a0ea11511e50039a7649619abfb4ba6d \
79+
--hash=sha256:93c36a08a6d158db44f2eb9ce3258e53f24a9a4a695325a689494f0fdbc71770 \
80+
--hash=sha256:95ce33057ba4d05df50b1f3cfefab22e351868a843b3b15a46c65836283670c9 \
81+
--hash=sha256:9849ce7267444c0a717c80c6900997de4f36e2815ce34ac560a3edb2d9a64cd2 \
82+
--hash=sha256:9d55ea7f7baf2aed61bf1872116cefc9d0c3693b45cae3916897ee27ef4b835e \
83+
--hash=sha256:a4f39b9bf6555fab9bfb536cf5fdd1c1c727e8d22312078702e9ff005354b37f \
84+
--hash=sha256:aec640bd94c4c85c0d11e2733bd13cbb10438fb004852996ec0efbc6cacdaf70 \
85+
--hash=sha256:aecbd7c5272c82e54d5b99d8435fd10915d1bc704b7df15e4d9ca8dc3902be61 \
86+
--hash=sha256:bda32ce212baa724e03c68771e5c69f39e584ea426bfe1a701cb01508ffc7035 \
87+
--hash=sha256:bdcf26c2dbc926b8a35ec8cbfad7eff1a8bd8239e12478caca83a1fc2c400dc2 \
88+
--hash=sha256:bdf40d2aaabd3913dec11840f0d0ebb1b93134f99af6a0a4fd88ffe924928ab4 \
89+
--hash=sha256:c205cac07d24a29840c163d6469f61069ce4b065518519216297fc2f261f8db9 \
90+
--hash=sha256:c3f5ae0309d9f888fd825c2e9d0241102fadaca81d888f26f845bc8c13c1e4ee \
91+
--hash=sha256:cd7c0bb22d4ff86d65ad61b5dd246812e8993fbc95b558553624c33e8b6903ea \
92+
--hash=sha256:d0f730a17cf4f343b2c7ad50cee3bd19e969e793d2be6ed911f43086460096e4 \
93+
--hash=sha256:da65e5fd3eea434ccb8984c3624bc234ddcc0d9f4c81864af611aaebcc08a50e \
94+
--hash=sha256:e12e29764a0e66a7a31e9b8bf1de5cc0423ea72979f45909acd4292de834ccd3

‎xla/tools/multihost_hlo_runner/BUILD‎

Lines changed: 32 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,17 @@
11
load("@bazel_skylib//rules:build_test.bzl", "build_test")
22
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
33
load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm")
4-
load("//xla:xla.default.bzl", "xla_cc_binary")
4+
load("//xla:xla.default.bzl", "xla_cc_binary", "xla_cc_library")
55
load("//xla/tests:build_defs.bzl", "xla_test")
66
load("//xla/tsl:tsl.bzl", "if_cuda_or_rocm", "if_google")
77
load("//xla/tsl:tsl.default.bzl", "tsl_pybind_extension")
88
load("//xla/tsl/platform:build_config_root.bzl", "tf_gpu_tests_tags")
9-
load("//xla/tsl/platform:rules_cc.bzl", "cc_library")
9+
load("//xla/tsl/platform:rules_cc.bzl", "cc_library", "cc_shared_library")
10+
load("@xla//xla:py_strict.bzl", "py_strict_binary")
11+
load(
12+
"@xla//third_party/py:py_import.bzl",
13+
"py_import",
14+
)
1015

1116
package(
1217
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -255,14 +260,17 @@ xla_test(
255260
],
256261
)
257262

263+
xla_cc_library(
264+
name = "runner_deps_so_hack",
265+
)
266+
258267
tsl_pybind_extension(
259268
name = "py_hlo_multihost_runner",
260269
srcs = ["python_hlo_runner.cc"],
261270
deps = [
271+
":runner_deps_so_hack",
262272
":create_client",
263273
":functional_hlo_runner",
264-
":hlo_input_output_format",
265-
":profiler_interface",
266274
"//xla:debug_options_flags",
267275
"//xla:status_macros",
268276
"//xla:xla_data_proto_cc",
@@ -278,7 +286,6 @@ tsl_pybind_extension(
278286
"//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config",
279287
"//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
280288
"//xla/python:logging",
281-
"//xla/service:cpu_plugin",
282289
"//xla/service:custom_call_target_registry",
283290
"//xla/service:hlo_module_util",
284291
"//xla/tsl/platform:statusor",
@@ -297,7 +304,6 @@ tsl_pybind_extension(
297304
"@tsl//tsl/platform:statusor",
298305
] + if_cuda_or_rocm([
299306
"//xla/service:gpu_plugin",
300-
"//xla/backends/profiler/gpu:cupti_tracer",
301307
"//xla/backends/profiler/gpu:device_tracer",
302308
]) + if_cuda([
303309
"//xla/stream_executor:cuda_platform",
@@ -309,3 +315,23 @@ tsl_pybind_extension(
309315
"//xla/stream_executor:rocm_platform",
310316
]),
311317
)
318+
319+
320+
# py_import(
321+
# name = "rocm_te",
322+
# wheel = ":transformer_engine_rocm-2.2.0-py3-none-manylinux_2_28_x86_64.whl",
323+
# wheel_deps = [],
324+
# deps = [],
325+
# )
326+
327+
328+
py_strict_binary(
329+
name = "python_hlo_runner_main",
330+
srcs = ["python_hlo_runner_main.py"],
331+
tags = ["gpu"],
332+
deps = [
333+
":py_hlo_multihost_runner",
334+
"@absl_py//absl/testing:absltest",
335+
# ":rocm_te",
336+
],
337+
)

‎xla/tools/multihost_hlo_runner/python_hlo_runner.cc‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,19 @@ limitations under the License.
3232
#include "xla/service/custom_call_target_registry.h"
3333
#include "xla/tools/multihost_hlo_runner/create_client.h"
3434
#include "xla/tools/multihost_hlo_runner/functional_hlo_runner.h"
35-
#include "xla/tools/multihost_hlo_runner/hlo_input_output_format.h"
3635
#include "xla/tsl/platform/statusor.h"
3736
#include "xla/xla_data.pb.h"
3837

38+
39+
// namespace tsl {
40+
// namespace profiler {
41+
// namespace internal {
42+
// std::atomic<int> g_trace_level(0);
43+
// std::atomic<uint64_t> g_trace_filter_bitmap(
44+
// std::numeric_limits<uint64_t>::max());
45+
// }}}
46+
47+
3948
namespace nb = ::nanobind;
4049

4150
namespace xla {
@@ -105,8 +114,8 @@ absl::StatusOr<FunctionalHloRunner::RunningOptions> RunningOptionsFromFlags(
105114
out.module_argument_mode = opts.hlo_argument_mode;
106115
out.module_output_mode = opts.output_mode;
107116
out.num_repeats = static_cast<size_t>(opts.num_repeats);
108-
out.num_repeats_with_profiler =
109-
static_cast<size_t>(opts.num_repeats_with_profiler);
117+
// out.num_repeats_with_profiler =
118+
// static_cast<size_t>(opts.num_repeats_with_profiler);
110119
out.log_input_output_mode =
111120
opts.log_output ? FunctionalHloRunner::LogOutputMode::kLogOutput
112121
: FunctionalHloRunner::LogOutputMode::kNotLogOutput;
@@ -142,6 +151,9 @@ RawCompileOptionsFromFlags(const PyHloRunnerConfig& opts) {
142151
return out;
143152
}
144153

154+
155+
156+
145157
absl::Status RunHloFiles(const std::vector<std::string>& hlo_files,
146158
const PyHloRunnerConfig& opts) {
147159
TF_ASSIGN_OR_RETURN(FunctionalHloRunner::PreprocessingOptions preproc_options,
@@ -407,9 +419,7 @@ NB_MODULE(py_hlo_multihost_runner, m) {
407419
.value("UnoptimizedSnapshotProtoBinary",
408420
InputFormat::kUnoptimizedSnapshotProtoBinary)
409421
.value("UnoptimizedSnapshotProtoText",
410-
InputFormat::kUnoptimizedSnapshotProtoText)
411-
.value("SerializedPjrtExecutable",
412-
InputFormat::kSerializedPjRtExecutable);
422+
InputFormat::kUnoptimizedSnapshotProtoText);
413423

414424
nb::enum_<FunctionalHloRunner::ModuleOutputMode>(m, "ModuleOutputMode")
415425
.value("ReturnOutputs",

‎xla/tools/multihost_hlo_runner/python_hlo_runner_main.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ==============================================================================
15+
16+
import os
17+
import pathlib
18+
import argparse
19+
20+
from transformer_engine import transformer_engine_jax
21+
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
22+
from xla.tools.multihost_hlo_runner import py_hlo_multihost_runner
23+
24+
25+
def _register_transformer_engine_custom_calls():
26+
for name, value in transformer_engine_jax.registrations().items():
27+
try:
28+
py_hlo_multihost_runner.register_custom_call_target(
29+
name, value, platform="ROCM", api_version=1
30+
)
31+
except:
32+
pass
33+
34+
35+
def main():
36+
parser = argparse.ArgumentParser(description="Run the specified hlo_file.")
37+
parser.add_argument("hlo_file", help="Path to the input file")
38+
parser.add_argument("-o", help="Path to output literal file")
39+
parser.add_argument(
40+
"--num_repeats",
41+
type=int,
42+
default=1,
43+
help="Number of times repeat the run (default: 1)",
44+
)
45+
46+
args = parser.parse_args()
47+
48+
_register_transformer_engine_custom_calls()
49+
50+
config = py_hlo_multihost_runner.PyHloRunnerConfig()
51+
config.input_format = (
52+
py_hlo_multihost_runner.InputFormat.SnapshotProtoBinary
53+
if ".pb" in args.hlo_file
54+
else py_hlo_multihost_runner.InputFormat.Text
55+
)
56+
config.hlo_argument_mode = (
57+
py_hlo_multihost_runner.ModuleArgumentMode.UseRandomInputs
58+
)
59+
config.log_output = True
60+
config.num_nodes = 1
61+
config.num_replicas = 1
62+
config.num_repeats = args.num_repeats
63+
config.output_mode = py_hlo_multihost_runner.ModuleOutputMode.ReturnDevice0Outputs
64+
if args.o:
65+
config.dump_output_literal_to = args.o
66+
67+
os.environ["NVTE_FUSED_ATTN"] = "1"
68+
os.environ["NVTE_FUSED_ATTN_CK"] = "1"
69+
os.environ["NVTE_CK_USES_FWD_V3"] = "1"
70+
os.environ["NVTE_CK_USES_BWD_V3"] = "1"
71+
os.environ["NVTE_CK_IS_V3_ATOMIC_FP32"] = "1"
72+
os.environ["NVTE_CK_HOW_V3_BF16_CVT"] = "1"
73+
74+
py_hlo_multihost_runner.RunHloFiles([args.hlo_file], config)
75+
print("DONE")
76+
77+
78+
if __name__ == "__main__":
79+
main()

‎xla/tsl/profiler/backends/cpu/traceme_recorder.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ namespace internal {
4545
// DLL imported variables cannot be initialized on Windows. This file is
4646
// included only on DLL exports.
4747
DECL_DLL_EXPORT std::atomic<int> g_trace_level(
48-
TraceMeRecorder::kTracingDisabled);
48+
TraceMeRecorder::kTracingDisabled);
4949
DECL_DLL_EXPORT std::atomic<uint64_t> g_trace_filter_bitmap(
5050
std::numeric_limits<uint64_t>::max());
5151

‎xla/tsl/tsl.bzl‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -726,7 +726,6 @@ def tsl_pybind_extension_opensource(
726726
compatible_with = compatible_with,
727727
deprecation = deprecation,
728728
features = features + ["-use_header_modules"],
729-
licenses = licenses,
730729
restricted_to = restricted_to,
731730
shared_lib_name = so_file,
732731
testonly = testonly,

‎xla/xla.default.bzl‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,9 @@ _XLA_SHARED_OBJECT_SENSITIVE_DEPS = if_static(extra_deps = [], otherwise = [
6969
def xla_cc_binary(deps = [], copts = tsl_copts(), **kwargs):
7070
native.cc_binary(deps = deps + _XLA_SHARED_OBJECT_SENSITIVE_DEPS, copts = copts, **kwargs)
7171

72+
def xla_cc_library(deps = [], copts = tsl_copts(), **kwargs):
73+
native.cc_library(deps = deps + _XLA_SHARED_OBJECT_SENSITIVE_DEPS, copts = copts, **kwargs)
74+
7275
def xla_cc_test(
7376
name,
7477
deps = [],

0 commit comments

Comments
 (0)