OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 12:20:11 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Enabling MLIR generated kernels by default on ROCm

Browse Source
This commit is contained in:
Deven Desai 2021-04-06 14:49:22 +00:00
parent a0ef5ae4e0
commit 22a512f570
2 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.bazelrc
View File

@ -248,8 +248,6 @@ build:tensorrt --repo_env TF_NEED_TENSORRT=1
build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain
build:rocm --define=using_rocm=true --define=using_rocm_hipcc=true
build:rocm --repo_env TF_NEED_ROCM=1
# Generated kernels are not yet supported on ROCm.
build:rocm --//tensorflow/core/kernels/mlir_generated:enable_gpu=false
# Options extracted from configure script
build:numa --define=with_numa_support=true

3
tensorflow/core/kernels/mlir_generated/build_defs.bzl
View File

@ -271,6 +271,7 @@ def _gen_kernel_library(
)
# We have to use a sh_test instead of build_test because it doesn't properly find the dependent targets.
gpu_arch_option = "sm_70,compute_75" if cuda_gpu_architectures() else ",".join(rocm_gpu_architectures())
native.sh_test(
name = "{op}_{platform}_{type}_{output_type}_gen_test".format(
op = op,
@ -288,7 +289,7 @@ def _gen_kernel_library(
type = type,
output_type = output_type,
),
"--cpu_codegen=true" if enable_cpu else "--arch=sm_70,compute_75",
"--cpu_codegen=true" if enable_cpu else "--arch={}".format(gpu_arch_option),
],
size = "medium",
data = [