From 24ee47b24613b184c7e3b74e1b4aa03e49d703b4 Mon Sep 17 00:00:00 2001 From: Deven Desai Date: Tue, 6 Apr 2021 14:49:22 +0000 Subject: [PATCH] Enabling MLIR generated kernels by default on ROCm --- .bazelrc | 2 -- tensorflow/core/kernels/mlir_generated/build_defs.bzl | 3 ++- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/.bazelrc b/.bazelrc index 9d57d2c2950..fcef170dded 100644 --- a/.bazelrc +++ b/.bazelrc @@ -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 diff --git a/tensorflow/core/kernels/mlir_generated/build_defs.bzl b/tensorflow/core/kernels/mlir_generated/build_defs.bzl index 530e61a7b5a..f03d2b1670d 100644 --- a/tensorflow/core/kernels/mlir_generated/build_defs.bzl +++ b/tensorflow/core/kernels/mlir_generated/build_defs.bzl @@ -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 = [