Skip to content

Commit 342b444

Browse files
authored
decompose disc-compiler with org tensorflow (#972)
decompose disc_compiler from tf_community
1 parent c681ec4 commit 342b444

File tree

400 files changed

+1503
-1562
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

400 files changed

+1503
-1562
lines changed

pytorch_blade/.bazelrc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
startup --host_jvm_args=-Djdk.http.auth.tunneling.disabledSchemes=
22

3-
try-import ../tf_community/.bazelrc
4-
try-import ../tf_community/.bazelrc.user
3+
try-import ../tao_compiler/.bazelrc
4+
try-import ../tao_compiler/.bazelrc.user
55

66
build --disk_cache=~/.cache
77
build --define is_torch_disc=true

pytorch_blade/WORKSPACE

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,11 @@ local_repository(
112112
path = "../tf_community",
113113
)
114114

115+
local_repository(
116+
name = "org_disc_compiler",
117+
path = "../tao_compiler",
118+
)
119+
115120
load("@org_third_party//bazel:blade_disc_dnn_workspace.bzl", "blade_disc_dnn_workspace")
116121

117122
blade_disc_dnn_workspace()

pytorch_blade/bazel/torch/torch.BUILD.tpl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ cc_library(
1515
"lib/libtorch_cpu.so",
1616
"lib/libtorch_global_deps.so",
1717
] + glob(
18-
["lib/libtorch_cuda.so","lib/libtensorpipe.so",]),
18+
["lib/libtorch_cuda_cu.so", "lib/libtorch_cuda_cpp.so", "lib/libgomp*.so.1"]),
1919
hdrs = glob(
2020
[
2121
"include/torch/**/*.h",

pytorch_blade/bazel_build.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ def pybind11_cflags(self):
6464
def __init__(self, *args, **kwargs):
6565
super().__init__(*args, **kwargs)
6666
self.targets = [
67-
"@org_tensorflow//tensorflow/compiler/mlir/xla/ral:libral_base_context.so",
67+
"@org_disc_compiler//mlir/xla/ral:libral_base_context.so",
6868
"//pytorch_blade:libtorch_blade.so",
6969
"//pytorch_blade:_torch_blade.so",
7070
"//tests/mhlo/torch-mlir-opt:torch-mlir-opt",
@@ -185,7 +185,7 @@ def __init__(self, *args, **kwargs):
185185
self.shell_setting = "set -e; set -o pipefail; "
186186
# Workaround: this venv ensure that $(/usr/bin/env python) is evaluated to python3
187187
venv.create(".bazel_pyenv", clear=True)
188-
self.build_cmd = "source .bazel_pyenv/bin/activate; bazel build"
188+
self.build_cmd = "source .bazel_pyenv/bin/activate; bazel build --verbose_failures"
189189
self.test_cmd = "source .bazel_pyenv/bin/activate; bazel test"
190190
if running_on_ci():
191191
self.test_cmd += " --test_output=errors"
@@ -212,7 +212,7 @@ def run(self, extdir=None, srcdir=None, build_temp=None):
212212
+ [
213213
"--compilation_mode=opt",
214214
"--define is_torch_disc=false", # still use mhlo within TF to build compiler main
215-
"@org_tensorflow//tensorflow/compiler/mlir/disc:disc_compiler_main",
215+
"@org_disc_compiler//mlir/disc:disc_compiler_main",
216216
]
217217
)
218218
subprocess.check_call(

pytorch_blade/distribution.cfg

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
pytorch_blade/_torch_blade.so
22
pytorch_blade/libtorch_blade.so
3-
external/org_tensorflow/tensorflow/compiler/mlir/xla/ral/libral_base_context.so
4-
external/org_tensorflow/tensorflow/compiler/mlir/disc/disc_compiler_main
3+
external/org_disc_compiler/mlir/xla/ral/libral_base_context.so
4+
external/org_disc_compiler/mlir/disc/disc_compiler_main

pytorch_blade/pytorch_blade/compiler/mlir/converters/mhlo_conversion.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
#include <mlir/Dialect/Tensor/IR/Tensor.h>
1717

1818
#include "llvm/Support/SourceMgr.h"
19-
#include "tensorflow/compiler/mlir/disc/IR/hlo_disc_ops.h"
19+
#include "mlir/disc/IR/hlo_disc_ops.h"
2020

2121
#include "pytorch_blade/common_utils/logging.h"
2222
#include "pytorch_blade/common_utils/utils.h"

pytorch_blade/pytorch_blade/compiler/mlir/runtime/BUILD

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ cc_library(
1111
deps = [
1212
"//pytorch_blade/common_utils:torch_blade_common",
1313
"//pytorch_blade/compiler/jit:torch_blade_jit",
14-
"@org_tensorflow//tensorflow/compiler/mlir/xla/ral:ral_base_context_lib",
14+
"@org_disc_compiler//mlir/xla/ral:ral_base_context_lib",
1515
"@local_org_torch//:ATen",
1616
"@local_org_torch//:libtorch",
1717
] + if_cuda_is_configured([

pytorch_blade/pytorch_blade/compiler/mlir/runtime/ral_context.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
#endif // TORCH_BLADE_USE_ROCM
2929
#endif // TORCH_BLADE_BUILD_WITH_CUDA
3030

31-
#include "tensorflow/compiler/mlir/xla/ral/ral_api.h"
31+
#include "mlir/xla/ral/ral_api.h"
3232

3333
#include "pytorch_blade/common_utils/utils.h"
3434

pytorch_blade/pytorch_blade/compiler/mlir/runtime/ral_context.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,11 @@ using CUDAStream = ::c10::hip::HIPStream;
3030
#else // TORCH_BLADE_USE_ROCM
3131
#include <c10/cuda/CUDAStream.h>
3232
#endif // TORCH_BLADE_USE_ROCM
33-
#include "tensorflow/compiler/mlir/xla/ral/context/base/cuda/cuda_context_impl.h"
33+
#include "mlir/xla/ral/context/base/cuda/cuda_context_impl.h"
3434
#endif // TORCH_BLADE_BUILD_WITH_CUDA
3535
// TODO(disc): figure out why the bazel does not trigger re-compile this file
3636
// after we update ral.
37-
#include "tensorflow/compiler/mlir/xla/ral/context/base/cpu/cpu_context_impl.h"
37+
#include "mlir/xla/ral/context/base/cpu/cpu_context_impl.h"
3838

3939
#include "pytorch_blade/common_utils/macros.h"
4040
#include "pytorch_blade/common_utils/tempfs.h"

pytorch_blade/pytorch_blade/torch-mlir/BUILD

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,8 @@ cc_library(
6868
":DiscTorchMLIRUtils",
6969
":TorchMLIRConversionMhloPassesIncGen",
7070
"@mlir-hlo//:mlir_hlo",
71-
"@org_tensorflow//tensorflow/compiler/mlir/disc:mhlo_disc",
72-
"@org_tensorflow//tensorflow/compiler/mlir/disc:disc_pdl_utils",
71+
"@org_disc_compiler//mlir/disc:mhlo_disc",
72+
"@org_disc_compiler//mlir/disc:disc_pdl_utils",
7373
"@llvm-project//mlir:Dialect",
7474
"@torch-mlir//:TorchMLIRTorchBackendTypeConversion",
7575
"@torch-mlir//:TorchMLIRTorchConversionDialect",

pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
#include "mlir/Pass/Pass.h"
2626
#include "mlir/Pass/PassManager.h"
2727
#include "mlir/Transforms/Passes.h"
28-
#include "tensorflow/compiler/mlir/disc/IR/hlo_disc_ops.h"
28+
#include "mlir/disc/IR/hlo_disc_ops.h"
2929
#include "torch-mlir/Dialect/Torch/Transforms/Passes.h"
3030

3131
namespace mlir {

pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/ApplyDiscPdlPatterns.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
#include "mlir/Interfaces/SideEffectInterfaces.h"
1515
#include "mlir/Transforms/DialectConversion.h"
1616
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
17-
#include "tensorflow/compiler/mlir/disc/transforms/disc_pdl_utils.h"
17+
#include "mlir/disc/transforms/disc_pdl_utils.h"
1818
#include "torch-mlir/Conversion/MhloPasses.h"
1919
#include "torch-mlir/Conversion/TorchToMhlo/TorchToMhlo.h"
2020
#include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"

pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616
#include "mlir/Dialect/Tensor/IR/Tensor.h"
1717
#include "mlir/IR/Builders.h"
1818
#include "mlir/Transforms/DialectConversion.h"
19+
#include "mlir/disc/IR/hlo_disc_ops.h"
1920
#include "stablehlo/dialect/ChloOps.h"
20-
#include "tensorflow/compiler/mlir/disc/IR/hlo_disc_ops.h"
2121
#include "torch-mlir/Conversion/MhloPasses.h"
2222
#include "torch-mlir/Conversion/TorchToMhlo/TorchToMhlo.h"
2323
#include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"

pytorch_blade/scripts/build_pytorch_blade.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ function ci_build() {
5353
python3 ../scripts/python/common_setup.py --cpu_only $COMMON_SETUP_ARGS
5454
fi
5555
TORCH_LIB=$(python -c 'import torch; import os; print(os.path.dirname(os.path.abspath(torch.__file__)) + "/lib/")') \
56-
export LD_LIBRARY_PATH=$TORCH_LIB:$LD_LIBRARY_PATH \
56+
export LD_LIBRARY_PATH=$TORCH_LIB:$LD_LIBRARY_PATH
5757

5858
# DEBUG=1 will trigger debug mode compilation
5959
DEBUG=1 python3 setup.py cpp_test 2>&1 | tee -a cpp_test.out;

pytorch_blade/tests/torch-disc-pdll/BUILD

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ cc_library(
1414
"@mlir-hlo//:mlir_hlo",
1515
"@llvm-project//mlir:IR",
1616
"@llvm-project//llvm:Support",
17-
"@org_tensorflow//tensorflow/compiler/mlir/disc:disc_pdl_utils",
17+
"@org_disc_compiler//mlir/disc:disc_pdl_utils",
1818
"@torch-mlir//:TorchMLIRTorchDialect",
1919
]
2020
)

pytorch_blade/tests/torch-disc-pdll/torch-disc-pdll.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
#include "mlir/Parser/Parser.h"
2020
#include "mlir/Support/FileUtilities.h"
2121
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
22-
#include "tensorflow/compiler/mlir/disc/IR/hlo_disc_ops.h"
22+
#include "mlir/disc/IR/hlo_disc_ops.h"
2323
#include "torch-mlir/InitAll.h"
2424

2525
#include <string>

pytorch_blade/tests/torch-disc-pdll/utils.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
#ifndef TORCH_DISC_PDLL_UTILS_H_
1313
#define TORCH_DISC_PDLL_UTILS_H_
1414

15-
#include "tensorflow/compiler/mlir/disc/transforms/disc_pdl_utils.h"
15+
#include "mlir/disc/transforms/disc_pdl_utils.h"
1616

1717
namespace mlir {
1818
class PDLPatternModule;

scripts/ci/build_and_test.sh

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ fi
4242

4343
# copy libtao_ops.so and tao_compiler_main to blade-disc-tf
4444
cp tao/bazel-bin/libtao_ops.so ${BLADE_DISC_DIR}
45-
cp tf_community/bazel-bin/tensorflow/compiler/decoupling/tao_compiler_main ${BLADE_DISC_DIR}
45+
cp tao_compiler/bazel-bin/decoupling/tao_compiler_main ${BLADE_DISC_DIR}
4646

4747
if [[ -n "$ROCM" ]] || [[ -n "$DCU" ]]; then
4848
# TODO: skip the following stages if rocm build
@@ -56,7 +56,7 @@ fi
5656
# copy Python wheel package to build folder
5757
mkdir -p build && \
5858
cp tao/dist/blade_disc*.whl ./build
59-
cp tf_community/bazel-bin/tensorflow/compiler/mlir/disc/tools/disc-replay/disc-replay-main ./build/
59+
#cp tao_compiler/bazel-bin/mlir/disc/tools/disc-replay/disc-replay-main ./build/
6060

6161
# test example models
6262
if [[ -z "$ROCM" ]] && [[ -z "$DCU" ]]; then

scripts/python/common_setup.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -611,10 +611,10 @@ def symlink_disc_files(args):
611611
os.makedirs(dst_folder)
612612
execute("rm -rf {0} && ln -s {1} {0}".format(link_in_tf, src_file))
613613

614-
logger.info("linking ./tao to tf_community/tao")
614+
logger.info("linking ./tao to tao_compiler/tao")
615615
execute(
616616
"rm -rf {0} && ln -s {1} {0}".format(
617-
os.path.join(get_source_root_dir(), "tf_community", "tao"),
617+
os.path.join(get_source_root_dir(), "tao_compiler", "tao"),
618618
os.path.join(get_source_root_dir(), "tao")
619619
)
620620
)
@@ -753,10 +753,15 @@ def test_tao_compiler_add_flags_platform_alibaba_cached(root, flag):
753753
restore_args_from_cache(root, args)
754754
return test_tao_compiler_add_flags_platform_alibaba(root, args, flag)
755755

756+
def cleanup_building_env(root):
757+
with cwd(root):
758+
execute("rm -rf tf_community/.tf_configure.bazelrc")
759+
756760
if __name__ == "__main__":
757761
args = parse_args()
758762

759763
root = get_source_root_dir()
764+
cleanup_building_env(root)
760765
symlink_disc_files(args)
761766

762767
if args.enable_mkldnn:

scripts/python/tao_build.py

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -279,10 +279,10 @@ def configure(root, args):
279279

280280
@time_stage()
281281
def build_tao_compiler(root, args):
282-
BAZEL_BUILD_CMD = "bazel build --experimental_multi_threaded_digest --define framework_shared_object=false" + ci_build_flag()
283-
TARGET_TAO_COMPILER_MAIN = "//tensorflow/compiler/decoupling:tao_compiler_main"
284-
TARGET_DISC_OPT = "//tensorflow/compiler/mlir/disc:disc-opt"
285-
TARGET_DISC_REPLAY = "//tensorflow/compiler/mlir/disc/tools/disc-replay:disc-replay-main"
282+
BAZEL_BUILD_CMD = "bazel build --verbose_failures --experimental_multi_threaded_digest --define framework_shared_object=false" + ci_build_flag()
283+
TARGET_TAO_COMPILER_MAIN = "//decoupling:tao_compiler_main"
284+
TARGET_DISC_OPT = "//mlir/disc:disc-opt"
285+
TARGET_DISC_REPLAY = "//mlir/disc/tools/disc-replay:disc-replay-main"
286286

287287
targets = None
288288
if args.bazel_target is not None:
@@ -296,9 +296,9 @@ def bazel_build(target, flag=""):
296296
logger.info("Building bazel target: " + target)
297297
execute(" ".join([BAZEL_BUILD_CMD, flag, target]))
298298

299-
with cwd(tf_root_dir(root)), gcc_env(args.compiler_gcc):
299+
with cwd(os.path.join(root, "tao_compiler")), gcc_env(args.compiler_gcc):
300300
execute(
301-
"cp -f -p {}/tao*.proto tensorflow/compiler/decoupling/".format(
301+
"cp -f -p {}/tao*.proto decoupling/".format(
302302
tao_bridge_dir(root)
303303
)
304304
)
@@ -341,7 +341,7 @@ def bazel_build(target, flag=""):
341341
if not args.rocm and not args.dcu:
342342
bazel_build(TARGET_DISC_REPLAY, flag=flag)
343343
execute(
344-
"cp -f -p {}/tao/third_party/ptxas/10.2/ptxas ./bazel-bin/tensorflow/compiler/decoupling/".format(
344+
"cp -f -p {}/tao/third_party/ptxas/10.2/ptxas ./bazel-bin/decoupling/".format(
345345
root
346346
)
347347
)
@@ -350,7 +350,7 @@ def bazel_build(target, flag=""):
350350

351351
@time_stage()
352352
def test_tao_compiler(root, args):
353-
BAZEL_BUILD_CMD = "bazel build --experimental_multi_threaded_digest --define framework_shared_object=false --test_timeout=600 --javabase=@bazel_tools//tools/jdk:remote_jdk11"
353+
BAZEL_BUILD_CMD = "bazel build --verbose_failures --experimental_multi_threaded_digest --define framework_shared_object=false --test_timeout=600 --javabase=@bazel_tools//tools/jdk:remote_jdk11"
354354
BAZEL_TEST_CMD = "bazel test --experimental_multi_threaded_digest --define framework_shared_object=false --test_timeout=600 --javabase=@bazel_tools//tools/jdk:remote_jdk11"
355355
BAZEL_TEST_CMD += ci_build_flag()
356356
BAZEL_BUILD_CMD += ci_build_flag()
@@ -362,24 +362,24 @@ def test_tao_compiler(root, args):
362362
BAZEL_TEST_CMD += " --java_runtime_version=remotejdk_11"
363363
BAZEL_BUILD_CMD += " --java_runtime_version=remotejdk_11"
364364

365-
TARGET_DISC_IR_TEST = "//tensorflow/compiler/mlir/disc/IR/tests/..."
366-
TARGET_DISC_TRANSFORMS_TEST = "//tensorflow/compiler/mlir/disc/transforms/tests/..."
367-
TARGET_DISC_E2E_TEST = "//tensorflow/compiler/mlir/disc/tests/..."
365+
TARGET_DISC_IR_TEST = "//mlir/disc/IR/tests/..."
366+
TARGET_DISC_TRANSFORMS_TEST = "//mlir/disc/transforms/tests/..."
367+
TARGET_DISC_E2E_TEST = "//mlir/disc/tests/..."
368368
TARGET_DISC_RAL_TESTS = [
369-
"//tensorflow/compiler/mlir/xla/ral:ral_metadata_test"
369+
"//mlir/xla/ral:ral_metadata_test"
370370
]
371371
TARGET_DISC_PDLL_TESTS = [
372-
"//tensorflow/compiler/mlir/disc/tools/disc-pdll/tests/..."
372+
"//mlir/disc/tools/disc-pdll/tests/..."
373373
]
374374
TARGET_DISC_CUDA_SOURCE_TESTS = [
375-
"//tensorflow/compiler/mlir/disc/tools/disc-source-emitter/tests/..."
375+
"//mlir/disc/tools/disc-source-emitter/tests/..."
376376
]
377377
TARGET_DISC_TRANSFORM_DIALECT_TESTS = [
378-
"//tensorflow/compiler/mlir/disc/tools/disc-transform/transforms/tests/...",
379-
"//tensorflow/compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/...",
378+
"//mlir/disc/tools/disc-transform/transforms/tests/...",
379+
"//mlir/disc/tools/disc-transform/LinalgExt/tests/...",
380380
]
381381

382-
TARGET_DISC_REPLAY_TEST = "//tensorflow/compiler/mlir/disc/tools/disc-replay:disc-replay-test"
382+
TARGET_DISC_REPLAY_TEST = "//mlir/disc/tools/disc-replay:disc-replay-test"
383383

384384
targets = None
385385
if args.bazel_target is not None:
@@ -394,9 +394,9 @@ def bazel_test(target, flag=""):
394394
execute(" ".join([BAZEL_BUILD_CMD, flag, target]))
395395
execute(" ".join([BAZEL_TEST_CMD, flag + ' --test_env=TF_CPP_VMODULE=disc_compiler=1,disc_transform_legalize_to_loop=1 --test_env=TF_ENABLE_ONEDNN_OPTS=0' , target]))
396396

397-
with cwd(tf_root_dir(root)), gcc_env(args.compiler_gcc):
397+
with cwd(os.path.join(root, "tao_compiler")), gcc_env(args.compiler_gcc):
398398
execute(
399-
"cp -f -p {}/tao*.proto tensorflow/compiler/decoupling/".format(
399+
"cp -f -p {}/tao*.proto decoupling/".format(
400400
tao_bridge_dir(root)
401401
)
402402
)
@@ -414,8 +414,8 @@ def bazel_test(target, flag=""):
414414
TARGET_DISC_TRANSFORMS_TEST,
415415
TARGET_DISC_E2E_TEST,
416416
] + TARGET_DISC_RAL_TESTS \
417-
+ TARGET_DISC_PDLL_TESTS \
418-
+ TARGET_DISC_TRANSFORM_DIALECT_TESTS
417+
+ TARGET_DISC_TRANSFORM_DIALECT_TESTS \
418+
+ TARGET_DISC_PDLL_TESTS
419419
MLIR_TESTS = " ".join(mlir_test_list)
420420
bazel_test(MLIR_TESTS, flag=flag)
421421
else:
@@ -640,7 +640,7 @@ def make_package(root, args):
640640
generate_build_info(build_info_file)
641641

642642
F_TAO_COMPILER_MAIN = (
643-
"./tf_community/bazel-bin/tensorflow/compiler/decoupling/tao_compiler_main"
643+
"./tao_compiler/bazel-bin/decoupling/tao_compiler_main"
644644
)
645645
F_TAO_OPS_SO = "./tao/bazel-bin/libtao_ops.so"
646646
F_PTXAS = "./tao/third_party/ptxas/10.2/ptxas"

tao/BUILD

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ cc_library(
103103
),
104104
deps = [
105105
":version_header",
106-
"@org_tao_compiler//tensorflow/compiler/mlir/xla/ral:ral_bridge",
106+
"@org_tao_compiler//mlir/xla/ral:ral_bridge",
107107
] + if_platform_alibaba([
108108
"@blade_service_common//blade_service_common:blade_service_common_deps",
109109
]) + if_internal_serving(

tao/tao_bridge/executable.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,9 @@
1818
#include <unordered_map>
1919
#include <vector>
2020

21+
#include "mlir/xla/ral/context/tensorflow/tf_context_impl.h"
2122
#include "tao_bridge/common.h"
2223
#include "tao_bridge/tao_compilation_result.pb.h"
23-
#include "tensorflow/compiler/mlir/xla/ral/context/tensorflow/tf_context_impl.h"
2424
#include "tensorflow/core/framework/op_kernel.h"
2525
#include "tensorflow/stream_executor/device_memory.h"
2626

tao/tao_bridge/mlir/mlir_executable.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
#include "tao_bridge/kernels/disc_launch.h"
2020
#endif // PLATFORM_ALIBABA
2121

22-
#include "tensorflow/compiler/mlir/xla/ral/ral_api.h"
22+
#include "mlir/xla/ral/ral_api.h"
2323

2424
namespace tensorflow {
2525
namespace tao {

tao/tao_bridge/test/tao_ut_common.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,8 @@ def _locate_tao_compiler():
3636
"""
3737
if 'TAO_COMPILER_PATH' not in os.environ:
3838
file_dir = os.path.abspath(os.path.dirname(__file__))
39-
compiler = os.path.join(file_dir, os.pardir, os.pardir, os.pardir, "tf_community",
40-
"bazel-bin", "tensorflow", "compiler", "decoupling",
41-
"tao_compiler_main")
39+
compiler = os.path.join(file_dir, os.pardir, os.pardir, os.pardir, "tao_compiler",
40+
"bazel-bin", "decoupling", "tao_compiler_main")
4241
compiler = os.path.abspath(compiler)
4342
assert os.path.exists(compiler), \
4443
"Tao compiler not found at: " + compiler

tao_compiler/.bazelrc

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
startup --host_jvm_args=-Djdk.http.auth.tunneling.disabledSchemes=
2+
3+
try-import ../tf_community/.bazelrc
4+
try-import ../tf_community/.tf_configure.bazelrc
5+
try-import %workspace%/.bazelrc.user
6+
7+
build --experimental_ui_max_stdouterr_bytes=-1

0 commit comments

Comments
 (0)