Update configure.py to provide a way to enable MLIR generated gpu kernels on the ROCm platform.
MLIR generated gpu kernels will continue to be disabled by default.
This commit provides a way to enable them, if conifugre.py is run with the env `TF_ENABLE_MLIR_GENERATED_GPU_KERNELS=1`
diff --git a/configure.py b/configure.py
index 9524ead..dc9d127 100644
--- a/configure.py
+++ b/configure.py
@@ -1442,6 +1442,10 @@
write_action_env_to_bazelrc('ROCM_PATH', environ_cp.get('ROCM_PATH'))
write_action_env_to_bazelrc('ROCM_ROOT', environ_cp.get('ROCM_PATH'))
+ if ((environ_cp.get('TF_NEED_ROCM') == '1') and
+ (environ_cp.get('TF_ENABLE_MLIR_GENERATED_GPU_KERNELS') == '1')):
+ write_to_bazelrc('build:rocm --define tensorflow_enable_mlir_generated_gpu_kernels=1')
+
environ_cp['TF_NEED_CUDA'] = str(
int(get_var(environ_cp, 'TF_NEED_CUDA', 'CUDA', False)))
if (environ_cp.get('TF_NEED_CUDA') == '1' and