enable tensorflow-2.14 and 2.16 for spack built ROCm (#44095)

* initial commit to enable tensorflow-2.14 for spack built ROCm

* fix style errors

* modify hipcc patch

* updates for rocm 6.1

* updates for tf-rocm-enhanced 2.16

* fix styling

* fix styling

* add patch for 2.16

* add patch for 2.16

* Update var/spack/repos/builtin/packages/py-tensorflow/package.py

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

* update rocm enhanced version names

* changes for rocm-enhanced version name change

* fix styling

---------

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

6 files changed, 846 insertions, 34 deletions

diff --git a/var/spack/repos/builtin/packages/hipcc/0014-remove-compiler-rt-linkage-for-host.6.0.patch b/var/spack/repos/builtin/packages/hipcc/0014-remove-compiler-rt-linkage-for-host.6.0.patch
index e49f055204..74b90fafe0 100644
--- a/var/spack/repos/builtin/packages/hipcc/0014-remove-compiler-rt-linkage-for-host.6.0.patch
+++ b/var/spack/repos/builtin/packages/hipcc/0014-remove-compiler-rt-linkage-for-host.6.0.patch
@@ -1,25 +1,24 @@
 diff --git a/bin/hipcc.pl b/bin/hipcc.pl
-index 513a427..cd2d6ac 100755
+index 513a427..780dc5c 100755
 --- a/bin/hipcc.pl
 +++ b/bin/hipcc.pl
-@@ -160,11 +160,14 @@ if ($HIP_PLATFORM eq "amd") {
+@@ -160,11 +160,13 @@ if ($HIP_PLATFORM eq "amd") {
      if($isWindows) {
          $execExtension = ".exe";
      }
 -    $HIPCC=get_normalized_path("$HIP_CLANG_PATH/clang++" . $execExtension);
 +    # llvm_path is set inside the hip recipe
-+    $LLVM_PATH= $ENV{'LLVM_PATH'};
-+    $HIPCC="${LLVM_PATH}/bin/clang++" . $execExtension;
++    $HIP_CLANG_PATH= $ENV{'HIP_CLANG_PATH'};
++    $HIPCC="${HIP_CLANG_PATH}/clang++" . $execExtension;
 
      # If $HIPCC clang++ is not compiled, use clang instead
      if ( ! -e $HIPCC ) {
 -        $HIPCC=get_normalized_path("$HIP_CLANG_PATH/clang" . $execExtension);
-+        $LLVM_PATH= $ENV{'LLVM_PATH'};
-+        $HIPCC="${LLVM_PATH}/bin/clang" . $execExtension;
++        $HIPCC="${HIP_CLANG_PATH}/clang" . $execExtension;
          $HIPLDFLAGS = "--driver-mode=g++";
      }
      # to avoid using dk linker or MSVC linker
-@@ -484,7 +487,8 @@ if($HIP_PLATFORM eq "amd"){
+@@ -484,7 +486,8 @@ if($HIP_PLATFORM eq "amd"){
              $targetsStr = $ENV{HCC_AMDGPU_TARGET};
          } elsif (not $isWindows) {
              # Else try using rocm_agent_enumerator
diff --git a/var/spack/repos/builtin/packages/hipcc/0014-remove-compiler-rt-linkage-for-host.6.1.patch b/var/spack/repos/builtin/packages/hipcc/0014-remove-compiler-rt-linkage-for-host.6.1.patch
index 3b9c3aa228..5bd44f33b8 100644
--- a/var/spack/repos/builtin/packages/hipcc/0014-remove-compiler-rt-linkage-for-host.6.1.patch
+++ b/var/spack/repos/builtin/packages/hipcc/0014-remove-compiler-rt-linkage-for-host.6.1.patch
@@ -1,26 +1,24 @@
 diff --git a/amd/hipcc/bin/hipcc.pl b/amd/hipcc/bin/hipcc.pl
-index 513a427..ffeafaa 100755
+index 513a427..780dc5c 100755
 --- a/amd/hipcc/bin/hipcc.pl
 +++ b/amd/hipcc/bin/hipcc.pl
-@@ -160,11 +160,15 @@ if ($HIP_PLATFORM eq "amd") {
+@@ -160,11 +160,13 @@ if ($HIP_PLATFORM eq "amd") {
      if($isWindows) {
          $execExtension = ".exe";
      }
 -    $HIPCC=get_normalized_path("$HIP_CLANG_PATH/clang++" . $execExtension);
-+    # llvm_path is set inside the hip recipe
-+    $LLVM_PATH= $ENV{'LLVM_PATH'};
-+    $HIPCC="${LLVM_PATH}/bin/clang++" . $execExtension;
-+
- 
++    # hip_clang_path is set inside the hip recipe
++    $HIP_CLANG_PATH= $ENV{'HIP_CLANG_PATH'};
++    $HIPCC="${HIP_CLANG_PATH}/clang++" . $execExtension;
+
      # If $HIPCC clang++ is not compiled, use clang instead
      if ( ! -e $HIPCC ) {
 -        $HIPCC=get_normalized_path("$HIP_CLANG_PATH/clang" . $execExtension);
-+        $LLVM_PATH= $ENV{'LLVM_PATH'};
-+        $HIPCC="${LLVM_PATH}/bin/clang" . $execExtension;
++        $HIPCC="${HIP_CLANG_PATH}/clang" . $execExtension;
          $HIPLDFLAGS = "--driver-mode=g++";
      }
      # to avoid using dk linker or MSVC linker
-@@ -484,7 +488,8 @@ if($HIP_PLATFORM eq "amd"){
+@@ -484,7 +486,8 @@ if($HIP_PLATFORM eq "amd"){
              $targetsStr = $ENV{HCC_AMDGPU_TARGET};
          } elsif (not $isWindows) {
              # Else try using rocm_agent_enumerator
diff --git a/var/spack/repos/builtin/packages/py-tensorflow/Find_ROCm_Components_Individiually.2.16.patch b/var/spack/repos/builtin/packages/py-tensorflow/Find_ROCm_Components_Individiually.2.16.patch
new file mode 100644
index 0000000000..aab9b09c53
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-tensorflow/Find_ROCm_Components_Individiually.2.16.patch
@@ -0,0 +1,364 @@
+diff --git a/third_party/gpus/find_rocm_config.py b/third_party/gpus/find_rocm_config.py
+index cd64efe6..e38495e2 100644
+--- a/third_party/gpus/find_rocm_config.py
++++ b/third_party/gpus/find_rocm_config.py
+@@ -70,9 +70,10 @@ def _get_header_version(path, name):
+ def _find_rocm_config(rocm_install_path):
+
+   def rocm_version_numbers(path):
++    rocm_core_path = os.environ.get('ROCM_CORE_PATH')
+     possible_version_files = [
+-        "include/rocm-core/rocm_version.h",  # ROCm 5.2
+-        "include/rocm_version.h",  # ROCm 5.1 and prior
++        rocm_core_path + "/include/rocm-core/rocm_version.h",
++        rocm_core_path + "/include/rocm_version.h",  # ROCm 5.1 and prior
+     ]
+     version_file = None
+     for f in possible_version_files:
+@@ -101,9 +102,10 @@ def _find_rocm_config(rocm_install_path):
+ def _find_hipruntime_config(rocm_install_path):
+
+   def hipruntime_version_number(path):
++    hip_path = os.environ.get('HIP_PATH')
+     possible_version_files = [
+-        "include/hip/hip_version.h",  # ROCm 5.2
+-        "hip/include/hip/hip_version.h",  # ROCm 5.1 and prior
++        hip_path + "/include/hip/hip_version.h",
++        hip_path + "/hip/include/hip/hip_version.h",  # ROCm 5.1 and prior
+     ]
+     version_file = None
+     for f in possible_version_files:
+@@ -133,7 +135,9 @@ def _find_hipruntime_config(rocm_install_path):
+ def _find_miopen_config(rocm_install_path):
+
+   def miopen_version_numbers(path):
++    miopen_path = os.environ.get('MIOPEN_PATH')
+     possible_version_files = [
++        miopen_path + "/include/miopen/version.h",
+         "include/miopen/version.h",  # ROCm 5.2 and prior
+         "miopen/include/miopen/version.h",  # ROCm 5.1 and prior
+     ]
+@@ -164,7 +168,9 @@ def _find_miopen_config(rocm_install_path):
+ def _find_rocblas_config(rocm_install_path):
+
+   def rocblas_version_numbers(path):
++    rocblas_path = os.environ.get('ROCBLAS_PATH')
+     possible_version_files = [
++        rocblas_path + "/include/rocblas/internal/rocblas-version.h",
+         "include/rocblas/internal/rocblas-version.h",  # ROCm 5.2
+         "rocblas/include/internal/rocblas-version.h",  # ROCm 5.1 and prior
+     ]
+@@ -196,7 +202,9 @@ def _find_rocblas_config(rocm_install_path):
+ def _find_rocrand_config(rocm_install_path):
+
+   def rocrand_version_number(path):
++    rocrand_path = os.environ.get('ROCRAND_PATH')
+     possible_version_files = [
++        rocrand_path + "/include/rocrand/rocrand_version.h",
+         "include/rocrand/rocrand_version.h",  # ROCm 5.1
+         "rocrand/include/rocrand_version.h",  # ROCm 5.0 and prior
+     ]
+@@ -222,7 +230,9 @@ def _find_rocrand_config(rocm_install_path):
+ def _find_rocfft_config(rocm_install_path):
+
+   def rocfft_version_numbers(path):
++    rocfft_path = os.environ.get('ROCFFT_PATH')
+     possible_version_files = [
++        rocfft_path + "/include/rocfft/rocfft-version.h",
+         "include/rocfft/rocfft-version.h",  # ROCm 5.2
+         "rocfft/include/rocfft-version.h",  # ROCm 5.1 and prior
+     ]
+@@ -253,7 +263,9 @@ def _find_rocfft_config(rocm_install_path):
+ def _find_hipfft_config(rocm_install_path):
+
+   def hipfft_version_numbers(path):
++    hipfft_path = os.environ.get('HIPFFT_PATH')
+     possible_version_files = [
++        hipfft_path + "/include/hipfft/hipfft-version.h",
+         "include/hipfft/hipfft-version.h",  # ROCm 5.2
+         "hipfft/include/hipfft-version.h",  # ROCm 5.1 and prior
+     ]
+@@ -284,7 +296,9 @@ def _find_hipfft_config(rocm_install_path):
+ def _find_roctracer_config(rocm_install_path):
+
+   def roctracer_version_numbers(path):
++    roctracer_path = os.environ.get('ROCTRACER_PATH')
+     possible_version_files = [
++        roctracer_path + "/include/roctracer/roctracer.h",
+         "include/roctracer/roctracer.h",  # ROCm 5.2
+         "roctracer/include/roctracer.h",  # ROCm 5.1 and prior
+     ]
+@@ -316,7 +330,9 @@ def _find_roctracer_config(rocm_install_path):
+ def _find_hipsparse_config(rocm_install_path):
+
+   def hipsparse_version_numbers(path):
++    hipsparse_path = os.environ.get('HIPSPARSE_PATH')
+     possible_version_files = [
++        hipsparse_path + "/include/hipsparse/hipsparse-version.h",
+         "include/hipsparse/hipsparse-version.h",  # ROCm 5.2
+         "hipsparse/include/hipsparse-version.h",  # ROCm 5.1 and prior
+     ]
+@@ -346,7 +362,9 @@ def _find_hipsparse_config(rocm_install_path):
+ def _find_hipsolver_config(rocm_install_path):
+
+   def hipsolver_version_numbers(path):
++    hipsolver_path = os.environ.get('HIPSOLVER_PATH')
+     possible_version_files = [
++        hipsolver_path + "/include/hipsolver/internal/hipsolver-version.h",
+         "include/hipsolver/internal/hipsolver-version.h",  # ROCm 5.2
+         "hipsolver/include/internal/hipsolver-version.h",  # ROCm 5.1
+         "hipsolver/include/hipsolver-version.h",  # ROCm 5.0 and prior
+@@ -378,7 +396,9 @@ def _find_hipsolver_config(rocm_install_path):
+ def _find_rocsolver_config(rocm_install_path):
+
+   def rocsolver_version_numbers(path):
++    rocsolver_path = os.environ.get('ROCSOLVER_PATH')
+     possible_version_files = [
++        rocsolver_path + "/include/rocsolver/rocsolver-version.h",
+         "include/rocsolver/rocsolver-version.h",  # ROCm 5.2
+         "rocsolver/include/rocsolver-version.h",  # ROCm 5.1 and prior
+     ]
+diff --git a/third_party/gpus/rocm_configure.bzl b/third_party/gpus/rocm_configure.bzl
+index 6eb17486..3e9a68d3 100644
+--- a/third_party/gpus/rocm_configure.bzl
++++ b/third_party/gpus/rocm_configure.bzl
+@@ -30,7 +30,6 @@ load(
+     "which",
+ )
+
+-
+ _GCC_HOST_COMPILER_PATH = "GCC_HOST_COMPILER_PATH"
+ _GCC_HOST_COMPILER_PREFIX = "GCC_HOST_COMPILER_PREFIX"
+ _ROCM_TOOLKIT_PATH = "ROCM_PATH"
+@@ -184,17 +183,23 @@ def _rocm_include_path(repository_ctx, rocm_config, bash_bin):
+     """
+     inc_dirs = []
+
++    hsa_path = get_host_environ(repository_ctx, "HSA_ROCR_DEV_PATH")
++    hip_path = get_host_environ(repository_ctx, "HIP_PATH")
++    rocprim_path = get_host_environ(repository_ctx, "ROCPRIM_PATH")
++    rocsolver_path = get_host_environ(repository_ctx, "ROCSOLVER_PATH")
++    rocblas_path = get_host_environ(repository_ctx, "ROCBLAS_PATH")
++    llvm_path = get_host_environ(repository_ctx, "LLVM_PATH")
+     # Add HSA headers (needs to match $HSA_PATH)
+-    inc_dirs.append(rocm_config.rocm_toolkit_path + "/hsa/include")
++    inc_dirs.append(hsa_path + "/include/hsa")
+
+     # Add HIP headers (needs to match $HIP_PATH)
+     inc_dirs.append(rocm_config.rocm_toolkit_path + "/hip/include")
+     if int(rocm_config.rocm_version_number) >= 50200:
+         inc_dirs.append(rocm_config.rocm_toolkit_path + "/include")
+-        inc_dirs.append(rocm_config.rocm_toolkit_path + "/include/hip")
+-        inc_dirs.append(rocm_config.rocm_toolkit_path + "/include/rocprim")
+-        inc_dirs.append(rocm_config.rocm_toolkit_path + "/include/rocsolver")
+-        inc_dirs.append(rocm_config.rocm_toolkit_path + "/include/rocblas")
++        inc_dirs.append(hip_path + "/include/hip")
++        inc_dirs.append(rocprim_path + "/include/rocprim")
++        inc_dirs.append(rocsolver_path + "/include/rocsolver")
++        inc_dirs.append(rocblas_path + "/include/rocblas")
+
+     # Add HIP-Clang headers (realpath relative to compiler binary)
+     rocm_toolkit_path = realpath(repository_ctx, rocm_config.rocm_toolkit_path, bash_bin)
+@@ -203,7 +208,7 @@ def _rocm_include_path(repository_ctx, rocm_config, bash_bin):
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/14.0.0/include")
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/15.0.0/include")
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/16.0.0/include")
+-    inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/17.0.0/include")
++    inc_dirs.append(llvm_path + "/lib/clang/17.0.0/include/")
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/17/include")
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/18/include")
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/19/include")
+@@ -230,7 +235,8 @@ def _amdgpu_targets(repository_ctx, rocm_toolkit_path, bash_bin):
+     """Returns a list of strings representing AMDGPU targets."""
+     amdgpu_targets_str = get_host_environ(repository_ctx, _TF_ROCM_AMDGPU_TARGETS)
+     if not amdgpu_targets_str:
+-        cmd = "%s/bin/rocm_agent_enumerator" % rocm_toolkit_path
++        rocminfo_path = get_host_environ(repository_ctx, "ROCMINFO_PATH")
++        cmd = rocminfo_path + "/bin/rocm_agent_enumerator"
+         result = execute(repository_ctx, [bash_bin, "-c", cmd])
+         targets = [target for target in result.stdout.strip().split("\n") if target != "gfx000"]
+         targets = {x: None for x in targets}
+@@ -351,23 +357,33 @@ def _find_libs(repository_ctx, rocm_config, hipfft_or_rocfft, miopen_path, rccl_
+     Returns:
+       Map of library names to structs of filename and path
+     """
++    hip_path = get_host_environ(repository_ctx, "HIP_PATH")
++    rocblas_path = get_host_environ(repository_ctx, "ROCBLAS_PATH")
++    hipfft_path = get_host_environ(repository_ctx, "HIPFFT_PATH")
++    hiprand_path = get_host_environ(repository_ctx, "HIPRAND_PATH")
++    hipsparse_path = get_host_environ(repository_ctx, "HIPSPARSE_PATH")
++    roctracer_path = get_host_environ(repository_ctx, "ROCTRACER_PATH")
++    rocsolver_path = get_host_environ(repository_ctx, "ROCSOLVER_PATH")
++    hipsolver_path = get_host_environ(repository_ctx, "HIPSOLVER_PATH")
++    hipblas_path = get_host_environ(repository_ctx, "HIPBLAS_PATH")
++
+     libs_paths = [
+         (name, _rocm_lib_paths(repository_ctx, name, path))
+         for name, path in [
+-            ("amdhip64", rocm_config.rocm_toolkit_path),
+-            ("rocblas", rocm_config.rocm_toolkit_path),
+-            (hipfft_or_rocfft, rocm_config.rocm_toolkit_path),
+-            ("hiprand", rocm_config.rocm_toolkit_path),
+-            ("MIOpen", miopen_path),
++            ("amdhip64", hip_path + "/lib"),
++            ("rocblas", rocblas_path),
++            (hipfft_or_rocfft, hipfft_path),
++            ("hiprand", hiprand_path),
++           ("MIOpen", miopen_path),
+             ("rccl", rccl_path),
+-            ("hipsparse", rocm_config.rocm_toolkit_path),
+-            ("roctracer64", rocm_config.rocm_toolkit_path),
+-            ("rocsolver", rocm_config.rocm_toolkit_path),
++            ("hipsparse", hipsparse_path),
++            ("roctracer64", roctracer_path + "/lib"),
++            ("rocsolver", rocsolver_path),
+         ]
+     ]
+     if int(rocm_config.rocm_version_number) >= 40500:
+-        libs_paths.append(("hipsolver", _rocm_lib_paths(repository_ctx, "hipsolver", rocm_config.rocm_toolkit_path)))
+-        libs_paths.append(("hipblas", _rocm_lib_paths(repository_ctx, "hipblas", rocm_config.rocm_toolkit_path)))
++        libs_paths.append(("hipsolver", _rocm_lib_paths(repository_ctx, "hipsolver", hipsolver_path)))
++        libs_paths.append(("hipblas", _rocm_lib_paths(repository_ctx, "hipblas", hipblas_path)))
+
+     # hipblaslt may be absent even in versions of ROCm where it exists
+     # (it is not installed by default in some containers). Autodetect.
+@@ -560,12 +576,13 @@ def _create_local_rocm_repository(repository_ctx):
+     hipfft_or_rocfft = "rocfft" if rocm_version_number < 40100 else "hipfft"
+
+     # For ROCm 5.2 and above, find MIOpen and RCCL in the main rocm lib path
+-    miopen_path = rocm_config.rocm_toolkit_path + "/miopen" if rocm_version_number < 50200 else rocm_config.rocm_toolkit_path
+-    rccl_path = rocm_config.rocm_toolkit_path + "/rccl" if rocm_version_number < 50200 else rocm_config.rocm_toolkit_path
++    miopen_path = get_host_environ(repository_ctx, "MIOPEN_PATH") + "/miopen"
++    rccl_path = get_host_environ(repository_ctx, "RCCL_PATH") + "/rccl"
+
+     # Copy header and library files to execroot.
+     # rocm_toolkit_path
+     rocm_toolkit_path = rocm_config.rocm_toolkit_path
++    rocblas_path = get_host_environ(repository_ctx, "ROCBLAS_PATH")
+     copy_rules = [
+         make_copy_dir_rule(
+             repository_ctx,
+@@ -578,6 +595,53 @@ def _create_local_rocm_repository(repository_ctx):
+     # explicitly copy (into the local_config_rocm repo) the $ROCM_PATH/hiprand/include and
+     # $ROCM_PATH/rocrand/include dirs, only once the softlink to them in $ROCM_PATH/include
+     # dir has been removed. This removal will happen in a near-future ROCm release.
++    rocm_parts = [
++        "rocblas",
++       "rocprim",
++        "hipfft",
++        "hipsparse",
++        "rocsolver",
++        "hipsolver",
++        "hipblas",
++       "rccl",
++       "miopen",
++       "roctracer",
++    ]
++    rocm_includes = ""
++
++    for part in rocm_parts:
++        part_cap = part.upper().replace("-","_")
++        part_path = get_host_environ(repository_ctx, part_cap + "_PATH")
++        rocm_includes = rocm_includes + '":' + part + '-include",\n'
++        rocm_includes = rocm_includes + '":' + part + '-include-second",\n'
++        copy_rules.append(
++            make_copy_dir_rule(
++                repository_ctx,
++                name = part + "-include",
++                src_dir = part_path + "/include/" + part,
++                out_dir = "rocm/include",
++            ),
++        )
++        copy_rules.append(
++            make_copy_dir_rule(
++                repository_ctx,
++                name = part + "-include-second",
++                src_dir = part_path + "/include/" + part,
++                out_dir = "rocm/include/" + part,
++            ),
++        )
++    hipcub_path = get_host_environ(repository_ctx, "HIPCUB_PATH")
++    copy_rules.append(
++        make_copy_dir_rule(
++            repository_ctx,
++            name = "hipcub-include",
++            src_dir = hipcub_path + "/include/hipcub/",
++            out_dir = "rocm/include/hipcub/",
++        )
++    )
++    rocm_includes = rocm_includes + '":hipcub-include",\n'
++
++    hiprand_path = get_host_environ(repository_ctx, "HIPRAND_PATH")
+     hiprand_include = ""
+     hiprand_include_softlink = rocm_config.rocm_toolkit_path + "/include/hiprand"
+     softlink_exists = files_exist(repository_ctx, [hiprand_include_softlink], bash_bin)
+@@ -587,11 +651,12 @@ def _create_local_rocm_repository(repository_ctx):
+             make_copy_dir_rule(
+                 repository_ctx,
+                 name = "hiprand-include",
+-                src_dir = rocm_toolkit_path + "/hiprand/include",
++                src_dir = hiprand_path + "/include/hiprand",
+                 out_dir = "rocm/include/hiprand",
+             ),
+         )
+
++    rocrand_path = get_host_environ(repository_ctx, "ROCRAND_PATH")
+     rocrand_include = ""
+     rocrand_include_softlink = rocm_config.rocm_toolkit_path + "/include/rocrand"
+     softlink_exists = files_exist(repository_ctx, [rocrand_include_softlink], bash_bin)
+@@ -601,11 +666,10 @@ def _create_local_rocm_repository(repository_ctx):
+             make_copy_dir_rule(
+                 repository_ctx,
+                 name = "rocrand-include",
+-                src_dir = rocm_toolkit_path + "/rocrand/include",
++                src_dir = rocrand_path + "/include/rocrand",
+                 out_dir = "rocm/include/rocrand",
+             ),
+         )
+-
+     rocm_libs = _find_libs(repository_ctx, rocm_config, hipfft_or_rocfft, miopen_path, rccl_path, bash_bin)
+     rocm_lib_srcs = []
+     rocm_lib_outs = []
+@@ -620,7 +684,8 @@ def _create_local_rocm_repository(repository_ctx):
+         outs = rocm_lib_outs,
+     ))
+
+-    clang_offload_bundler_path = rocm_toolkit_path + "/llvm/bin/clang-offload-bundler"
++    llvm_amdgpu_path = get_host_environ(repository_ctx, "LLVM_PATH")
++    clang_offload_bundler_path = llvm_amdgpu_path + "/bin/clang-offload-bundler"
+
+     # copy files mentioned in third_party/gpus/rocm/BUILD
+     copy_rules.append(make_copy_files_rule(
+@@ -665,9 +730,7 @@ def _create_local_rocm_repository(repository_ctx):
+         "%{roctracer_lib}": rocm_libs["roctracer64"].file_name,
+         "%{rocsolver_lib}": rocm_libs["rocsolver"].file_name,
+         "%{copy_rules}": "\n".join(copy_rules),
+-        "%{rocm_headers}": ('":rocm-include",\n' +
+-                            hiprand_include +
+-                            rocrand_include),
++        "%{rocm_headers}": ('":rocm-include",\n' + rocrand_include + hiprand_include + rocm_includes),
+     }
+     if rocm_libs["hipblaslt"] != None:
+         repository_dict["%{hipblaslt_lib}"] = rocm_libs["hipblaslt"].file_name
+@@ -735,16 +798,18 @@ def _create_local_rocm_repository(repository_ctx):
+         tpl_paths["crosstool:hipcc_cc_toolchain_config.bzl"],
+     )
+
++    hip_path = get_host_environ(repository_ctx, "HIP_PATH")
++    hsa_rocr_dev_path = get_host_environ(repository_ctx, "HSA_ROCR_DEV_PATH")
+     repository_ctx.template(
+         "crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc",
+         tpl_paths["crosstool:clang/bin/crosstool_wrapper_driver_rocm"],
+         {
+             "%{cpu_compiler}": str(cc),
+-            "%{hipcc_path}": rocm_config.rocm_toolkit_path + "/bin/hipcc",
++            "%{hipcc_path}": hip_path + "/bin/hipcc",
+             "%{hipcc_env}": _hipcc_env(repository_ctx),
+-            "%{rocr_runtime_path}": rocm_config.rocm_toolkit_path + "/lib",
++            "%{rocr_runtime_path}": hsa_rocr_dev_path + "/lib",
+             "%{rocr_runtime_library}": "hsa-runtime64",
+-            "%{hip_runtime_path}": rocm_config.rocm_toolkit_path + "/lib",
++            "%{hip_runtime_path}": hip_path + "/lib",
+             "%{hip_runtime_library}": "amdhip64",
+             "%{crosstool_verbose}": _crosstool_verbose(repository_ctx),
+             "%{gcc_host_compiler_path}": str(cc),
diff --git a/var/spack/repos/builtin/packages/py-tensorflow/Find_ROCm_Components_Individiually.patch b/var/spack/repos/builtin/packages/py-tensorflow/Find_ROCm_Components_Individiually.patch
new file mode 100644
index 0000000000..58ac82b429
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-tensorflow/Find_ROCm_Components_Individiually.patch
@@ -0,0 +1,388 @@
+diff --git a/tensorflow/tsl/platform/default/rocm_rocdl_path.cc b/tensorflow/tsl/platform/default/rocm_rocdl_path.cc
+index dfd9721b..88215316 100644
+--- a/tensorflow/tsl/platform/default/rocm_rocdl_path.cc
++++ b/tensorflow/tsl/platform/default/rocm_rocdl_path.cc
+@@ -40,6 +40,9 @@ string RocmRoot() {
+ #endif
+ }
+
+-string RocdlRoot() { return io::JoinPath(RocmRoot(), "amdgcn/bitcode"); }
++string RocdlRoot() {
++  const char* llvm_path_env = std::getenv("LLVM_PATH");
++  return io::JoinPath(llvm_path_env, "amdgcn/bitcode");
++}
+
+ }  // namespace tsl
+diff --git a/third_party/gpus/find_rocm_config.py b/third_party/gpus/find_rocm_config.py
+index cd64efe6..e38495e2 100644
+--- a/third_party/gpus/find_rocm_config.py
++++ b/third_party/gpus/find_rocm_config.py
+@@ -70,9 +70,10 @@ def _get_header_version(path, name):
+ def _find_rocm_config(rocm_install_path):
+
+   def rocm_version_numbers(path):
++    rocm_core_path = os.environ.get('ROCM_CORE_PATH')
+     possible_version_files = [
+-        "include/rocm-core/rocm_version.h",  # ROCm 5.2
+-        "include/rocm_version.h",  # ROCm 5.1 and prior
++        rocm_core_path + "/include/rocm-core/rocm_version.h",
++        rocm_core_path + "/include/rocm_version.h",  # ROCm 5.1 and prior
+     ]
+     version_file = None
+     for f in possible_version_files:
+@@ -101,9 +102,10 @@ def _find_rocm_config(rocm_install_path):
+ def _find_hipruntime_config(rocm_install_path):
+
+   def hipruntime_version_number(path):
++    hip_path = os.environ.get('HIP_PATH')
+     possible_version_files = [
+-        "include/hip/hip_version.h",  # ROCm 5.2
+-        "hip/include/hip/hip_version.h",  # ROCm 5.1 and prior
++        hip_path + "/include/hip/hip_version.h",
++        hip_path + "/hip/include/hip/hip_version.h",  # ROCm 5.1 and prior
+     ]
+     version_file = None
+     for f in possible_version_files:
+@@ -133,7 +135,9 @@ def _find_hipruntime_config(rocm_install_path):
+ def _find_miopen_config(rocm_install_path):
+
+   def miopen_version_numbers(path):
++    miopen_path = os.environ.get('MIOPEN_PATH')
+     possible_version_files = [
++        miopen_path + "/include/miopen/version.h",
+         "include/miopen/version.h",  # ROCm 5.2 and prior
+         "miopen/include/miopen/version.h",  # ROCm 5.1 and prior
+     ]
+@@ -164,7 +168,9 @@ def _find_miopen_config(rocm_install_path):
+ def _find_rocblas_config(rocm_install_path):
+
+   def rocblas_version_numbers(path):
++    rocblas_path = os.environ.get('ROCBLAS_PATH')
+     possible_version_files = [
++        rocblas_path + "/include/rocblas/internal/rocblas-version.h",
+         "include/rocblas/internal/rocblas-version.h",  # ROCm 5.2
+         "rocblas/include/internal/rocblas-version.h",  # ROCm 5.1 and prior
+     ]
+@@ -196,7 +202,9 @@ def _find_rocblas_config(rocm_install_path):
+ def _find_rocrand_config(rocm_install_path):
+
+   def rocrand_version_number(path):
++    rocrand_path = os.environ.get('ROCRAND_PATH')
+     possible_version_files = [
++        rocrand_path + "/include/rocrand/rocrand_version.h",
+         "include/rocrand/rocrand_version.h",  # ROCm 5.1
+         "rocrand/include/rocrand_version.h",  # ROCm 5.0 and prior
+     ]
+@@ -222,7 +230,9 @@ def _find_rocrand_config(rocm_install_path):
+ def _find_rocfft_config(rocm_install_path):
+
+   def rocfft_version_numbers(path):
++    rocfft_path = os.environ.get('ROCFFT_PATH')
+     possible_version_files = [
++        rocfft_path + "/include/rocfft/rocfft-version.h",
+         "include/rocfft/rocfft-version.h",  # ROCm 5.2
+         "rocfft/include/rocfft-version.h",  # ROCm 5.1 and prior
+     ]
+@@ -253,7 +263,9 @@ def _find_rocfft_config(rocm_install_path):
+ def _find_hipfft_config(rocm_install_path):
+
+   def hipfft_version_numbers(path):
++    hipfft_path = os.environ.get('HIPFFT_PATH')
+     possible_version_files = [
++        hipfft_path + "/include/hipfft/hipfft-version.h",
+         "include/hipfft/hipfft-version.h",  # ROCm 5.2
+         "hipfft/include/hipfft-version.h",  # ROCm 5.1 and prior
+     ]
+@@ -284,7 +296,9 @@ def _find_hipfft_config(rocm_install_path):
+ def _find_roctracer_config(rocm_install_path):
+
+   def roctracer_version_numbers(path):
++    roctracer_path = os.environ.get('ROCTRACER_PATH')
+     possible_version_files = [
++        roctracer_path + "/include/roctracer/roctracer.h",
+         "include/roctracer/roctracer.h",  # ROCm 5.2
+         "roctracer/include/roctracer.h",  # ROCm 5.1 and prior
+     ]
+@@ -316,7 +330,9 @@ def _find_roctracer_config(rocm_install_path):
+ def _find_hipsparse_config(rocm_install_path):
+
+   def hipsparse_version_numbers(path):
++    hipsparse_path = os.environ.get('HIPSPARSE_PATH')
+     possible_version_files = [
++        hipsparse_path + "/include/hipsparse/hipsparse-version.h",
+         "include/hipsparse/hipsparse-version.h",  # ROCm 5.2
+         "hipsparse/include/hipsparse-version.h",  # ROCm 5.1 and prior
+     ]
+@@ -346,7 +362,9 @@ def _find_hipsparse_config(rocm_install_path):
+ def _find_hipsolver_config(rocm_install_path):
+
+   def hipsolver_version_numbers(path):
++    hipsolver_path = os.environ.get('HIPSOLVER_PATH')
+     possible_version_files = [
++        hipsolver_path + "/include/hipsolver/internal/hipsolver-version.h",
+         "include/hipsolver/internal/hipsolver-version.h",  # ROCm 5.2
+         "hipsolver/include/internal/hipsolver-version.h",  # ROCm 5.1
+         "hipsolver/include/hipsolver-version.h",  # ROCm 5.0 and prior
+@@ -378,7 +396,9 @@ def _find_hipsolver_config(rocm_install_path):
+ def _find_rocsolver_config(rocm_install_path):
+
+   def rocsolver_version_numbers(path):
++    rocsolver_path = os.environ.get('ROCSOLVER_PATH')
+     possible_version_files = [
++        rocsolver_path + "/include/rocsolver/rocsolver-version.h",
+         "include/rocsolver/rocsolver-version.h",  # ROCm 5.2
+         "rocsolver/include/rocsolver-version.h",  # ROCm 5.1 and prior
+     ]
+diff --git a/third_party/gpus/find_rocm_config.py.gz.base64 b/third_party/gpus/find_rocm_config.py.gz.base64
+index fee67ec5..6e437b26 100644
+--- a/third_party/gpus/find_rocm_config.py.gz.base64
++++ b/third_party/gpus/find_rocm_config.py.gz.base64
+@@ -1 +1 @@
+-eJztW21v2zgS/q5fQSgoam9dJSlugUMOOcCbZlHftUlhZ7tYtIVB27TNrSz6SCppUPS/7wxJyRIt2YqjvX6JgTa2NPNwOC8Ph7Z4RC7E+l7yxVKTVyevTsjNkpEblighf43FHemneimkikg/jskQxRQZMsXkLZtFwVFwRN7yKYizGUmTGZNEg35/Tafwx93pkQ9MKi4S8io6IR0UCN2tsPsvQLgXKVnRe5IITVLFAIIrMucxI+zrlK014QmZitU65jSZMnLH9dIM40DADPKHgxATTUGagvwaPs2LcoRqYzC+llqvz46P7+7uImqMjYRcHMdWUB2/HVxcXo0uX4LBRuW3JGZKEcn+l3IJU53cE7oGe6Z0AlbG9I4ISehCMrinBdp7J7nmyaJHlJjrOyoZoMy40pJPUl1yVmYdzLkoAO6iCQn7IzIYheSX/mgw6gHG74ObN9e/3ZDf+8Nh/+pmcDki10NycX31enAzuL6CT7+S/tUf5L+Dq9c9wsBVMAz7upZoPxjJ0Y0mdGTEWMmAubAGqTWb8jmfwrySRUoXjCzELZMJTIesmVxxhcFUYN4MUGK+4ppqc2VrUjjMeauvIAzD95InmIbXFysYfiKpvEdjyJJRHH8GIZpqITkzNpJbm32QUgIMRMeaWd4rzVZREGDCq6nkkGeKUQm5oIwr6uAxMVUZpQcRR69pFcDFFabAjGl0VWJczGVmhAFaW/tRfyqSOV+k0jgQ9ZSeiVRHxqo1xUQXGThmiIsNptlSinSxxCRhyS2XIlmxRJNbKrlJyg7Y/278vn/zphsFgzkUF9yL+cwbkju39Ox0rB8yA405TEoTasl0Kk3YCVwCB03FjJX9p+kXZueVxeC+YDEUDd7K7aq0OyrixUJ8scGwvrfxzGJiA2GqfUnl7CXaM4MYaqj7QKWTYh7MpViRCVXOqY4YNrbl9kYEfLUxEdwDrBTkgsZNUJbHYq2PpZiuQhRJkf4o2KIh7nOaxjifOGUBZmsQQM0JCeET2TuhsnfAC+4dZFIQBNOYQp1emBBdopc7l4YCIVTds4CA9QrFYBQyXjA9dsON0ZQxTq1jxGysimYWlYwwJJWmcVxQAltfZ1lrPZ2F3IVtRZySTR3UjHCCMJ6PSM5r7QNxPidh7uMQQyhU5LIBLakG3Mh8LGh/BvmjOnljomQ0NkNvCXU3rtq6V3QYLj5Ccc3GrozHSbqaMNlZ0T+F7BHwGP4Btemy6P/TE3iRn4gRIy/wM35CafhkxIvD2CTPxuigFT2S0BXLwjN08QA2X4PBQOkABOomYODFiyJHudBg9cTIQ3Cfi0isWYYcyhBWiATKBpj9PEz1/OU/w671/wptAx9KFpm3HRke2YHIM0VedD7NXnRD8sxY1zP4XaMHkTXyFoXYMgAcczFaAGWtO6ddd9N5CbimY+S6AbqOcqjSYgU8z0YOv3035WZXtU/J8wgmB8gd4yPywsGWXyExfQUugMgz4AXTWXz7DuvmpyTMIEw61ENwoEM7F8sIyIJgN1sw9C4UDY3DbhZLMHZmM95ybUXineFUUdjcKmeV6jgRHBvyTnHgxVzGst45+ZibGvJkGqczZgodWFDad5lGtIQwQ4mYAv45elWtVyN9mi1YQhq9z+b/oi1gypVImLmOuTZHF1dbnSdF4aJfrH8KnuXnPMsTSKrsNvsKLZLqbCF0z/JpecZtieaCEyCGL1nWYkyLohnedjoWUiQ0XsqWdjMg4tg+A9wA+Zrll5dY1Q7qmgogji/OK1mhqNBzNPrhcjiC5m/8rv+f62Fo3WZZ5gCMwVWOsXYs8FAMIOaLNw7DFXkFU+JcKy4j61QVRQV7ZyuPLTNQ/ObcHFYghGcHkDngfQ+8VcKOVir2JV/LNNF8xRqUfEHYs+DAugdA/Le/4FGwmdJT3e+s+zeD92RoY/g3lP8RbMJhmSms5dDg2sb7K246oc3MlkT0PBiTVV6P3C0F2G27T4sFSJ2CiCUJaEOwGXlBSnew9LtODXoNyfit3SJmDYwlFdsavpzSeJpCLwjWMQkcoLjZnUPzL0zTbHHuKKy6UHe4XAoDNkkVNiwKu/A1hfnDvhK6f0gH3KeqhzLg1tweToBbTihz12mpiTO4GKitoi8SUG2RAw3VE0Blh1okoK0xSzS04tjeNaAgJ/i4vqPEQRbxuMQm7rVhIo9UchCnvAvriZkaMdPzd4Nr8F2ZlEzfnDNT3jUXwbvdh5YdDPT+8uqxleejHNZ8eCiPbT9qqqO6ASnVXJEBKlHCs0KwHt2OlMb2dx+TmKpmGxAj2d4eBOGOcXMkExpnF17ubU42mhapIcITH+zsVJzz2u1SHkgVEKlf3vZHLWxTyjAH71RKMC1sVioLqHa/UihMb8tSgdMuYZRH9xlDQiU1Ywwj2c7mxaEde6h+lZd4wih4ADWKJ0/0sIceZOE3kt30sJMPytnQuBSH/avXWSmWi7AM6GqnkKJe7VSkJLTaNbm6t88uD+UXynyum9UJCra3sAKa+9NoNUXxsvbTEnpgjYDvfugK6qWS0Xz4AuqjoObD108PxWg+bvmsqpLa1XNTex4BbKO0vnZuxva//WvGCE6wJUawaO7PfkZw4mXtJ0Y4hBGs734oI1gTPtiL7w7jgzLGYWxQwnj/SC6oqY9qLihVnffV29/NBaWx/e5ASzoFhzVqEJxsez2CBdy829khOOEt9Sc2eGB/YN32ozfZN8P+xeWwhW22D1TYaB9tsjZ/3kow+5v4kt4yQl01Z75w7XuRVk4e0yzUVExtv1CqRq9lqMRqvWsoWeA3Dgp/AWn4q6GTba99sICbd42aCKe0BfPUShzYSlj3/ehuwlrx6IaiDHNwT1GCaaGtqCmd2s6iVJZec1GJ1Xp/UbLAJw0R3zbqMDayLZKGAdx8PZ9fasYembb3Jf8+lNNdKPuUn74E3OKgjRcfTDxtko6x4PGkU4I5nHSKMG2QTnXp1ZNOsax90qnCap90ihb4+5rGrLORbW9f4yo+f9foG1CPJvYpP7Uq+/Y5BzJGm6wBsRpdv/3Qyj7HBzr4B0UPqIWfFB/EHX5xetuc/wN3+BY47th6Etp/kp2SGZ/iaQI80CHmthKNIQnDIyhuPjyZi31HDSpONQR5tnslVUVeewrg+Sg/LLI5HmKerMn3vwY8f7xmewwbK6bwfAhE6Pvm40f7wKoWIv7CtZEOP2cPwpZOJWQaUbqeUc064wYPm3drtJo8tVqnu+9Rszq9vQ+m7FDc/fv0DsWd3867LCkFwiuVz+Tf5+QfJ6cnJy5Nqp25d5ga8/Z9Ybgjfvu+Pmg+uZ/3TG7f6r9jfvtVC0RiIBx9rCisj6Y0tbw/y1fdL+y+l53GSIgSUrNZZ5tqIqCylep08+XSHGvrhM/UGXmm8PRKZ4Nk7HfHTQv1j6e63PKo7lVkj8NFeLiTdcJPyeVweD08g1L+lBTOkigtOwDYzdWAGDQeewkCiMV4jMdVxmNyfk7C8RjnOB4bNrbTDf4CEAHDYg==
+\ No newline at end of file
++eJztW21v2zgS/q5fQSgoKl8cJSlugUMOOcCbpqjv2iSws10smsKQbdrmVhZ9pJQ0KPrfb4akZIqW5NfbfkmA1rY083A4L4+GknhErvjiWbDpLCVvzt6ckfsZJfc0kVy8i/kT6WTpjAsZkk4ckx6KSdKjkopHOg69I++IfGAjEKdjkiVjKkgK+p1FNIIPc6ZNPlEhGU/Im/CMBCjgm1N+65+A8MwzMo+eScJTkkkKEEySCYspod9GdJESlpARny9iFiUjSp5YOlPDGBAwg/xhIPgwjUA6AvkF/JrYciRKlcH4N0vTxcXp6dPTUxgpY0MupqexFpSnH7pX1zf96xMwWKn8lsRUSiLofzMmYKrDZxItwJ5RNAQr4+iJcEGiqaBwLuVo75NgKUumbSL5JH2KBAWUMZOpYMMsLTkrtw7mbAuAu6KE+J0+6fZ98mun3+23AeP37v3729/uye+dXq9zc9+97pPbHrm6vXnbve/e3sCvd6Rz8wf5T/fmbZtQcBUMQ78tBNoPRjJ0owod6VNaMmDCtUFyQUdswkYwr2SaRVNKpvyRigSmQxZUzJnEYEowbwwoMZuzNErVkZVJ4TCXB/3zfN+/EyzBNLy9msPwQxGJZzSGzGiE448hRKOUC0aVjeRRZx+kFAcD0bFqls8ypfPQ8zDh5UgwyDNJIwG5IJUr6uAxMWUZpQ0RR6+l0oODc0yBMU3RVYlyMRO5EQpooe1H/RFPJmyaCeVA1JPpmGdpqKxaRJjoPAfHDDGxwTSbCZ5NZ5gkNHlkgidzmqTkMRJMJWUA9n8c3HXu37dCrzuB4oJzMRs7QzLjlraejvZDbqAyhwqhQi1omgkVdgKHwEEjPqZl/6XRV6rnlcfg2bIYigZPFXZV2h3aeDHnX3UwtO91PPOY6ECoap9FYnyC9owhhinUvSezoZ0HE8HnZBhJ41RDDEvbCntDAr5amgjuAVbyCkHlJijLU75ITwUfzX0UyZD+IrAlhbhPoizG+cQZ9TBbPQ9qjgsIH8+/cZl/A14w3yCTPM8bxRHU6ZUK0TV6ObhWFAihal14BKyXKAajkMGUpgMz3ABNGeDUAiWmY2WbaSspYUiqNIpjSwlsfZtnrfZ0HnITtjkxSjp1UDPECcJ4LiK5rLUPxNmE+IWPfQwhl6HJBrSkGnAp89nS/gLyR3XyykRBo1gNvSLUWrpq5ZztMLz4cMlSOjBlPEiy+ZCKYB79yUWbgMfwA9RGM9v/52fwR/5GlBg5xt/4C6XhlxK3h9FJno8RoBVtkkRzmoenZ+IBbL4Ag4HSAQjUVcDAi1c2R5nQYPXEyENwnvGQL2iO7AsfrhAJlA0w+6WfpZOTf/gt7f852gY+FDRUXwPhH+mByCtJjoOH8XHLJ6+UdW2F31J6EFklr1GILgPAUQfDKVDWIjhvmZPGS8A1gZJreei6iEGV2hXwOh/Z//5DlZu+qj0kr0OYHCAHykfk2MCW/3yi+gq8ACLPgBdUZ/H9B1w3HxI/h1DpUA/BgA71XDQjIAuC3XRK0btQNFHst/JYgrFjnfGaaysS7wKnisLqVDmrZGBEikIYcUFXqyCEnAleq0q4uu1dq3J4rX0L2SoZsGmBrLnyknwuJugAHwNZsGQUZ2OqCOMET53a1oUzv72ptq2D5anI45fwPL/8caGQvqj/bRvBxBueUHUcM3eCAaueTZFi1kG39P/kLM/2SZ51kKL5afoNGi4ZrCC0LoqJOsatiBaCQ6CZr3kNYIbYojneanJbCecrL+WNghoQcXTXAm6A7C+ytdIjLVVAxNDNZSWp2Aptw8Kfrnt96B0HHzv/vu352k+apHbA6N4UGAtDIttiQCJfvTcYhiMqiBbnWnEYSauqpirI3yvyGOMBit9NNPwKBP9ih2sB4P3wnIuMHq3EFTO2EFmSsjndgDEsYccCK3VBqI4w3nfvtqSKAswucziI/yrZoaSAgg1KL/TQRA8+RIv0dLg3owaLUJpZ4giW+nAxszoGaKN1e/8Nl7bQzOYXXnQ1po0p0DZ5mnEwVPe4GguQAktEcwk0O9jyHJPSGWSIllGDjkZQ9qgXonmbpLlHN6AnoygeZdBxgnVUAFVIpu4BwBKDq9Zc4zxFcG2H9MaLMldgw0xiWySx119EMH9YvcIaA+KPq2G5LVGuzG17nlxxQpnizkutosLFQK1wg81TtVwAbFXPE5V9sM1TK2OW2GrOsIncgKmMYEN3YyRqmOpj9/bu+mZLsrIhbb7Sx0+r6Mqvl7HY6Y3DTkrTaGwE8EJvTd0PRht8VyY51fEXTFf0+zZ4q7VtKZu02rOaXZTd+h4HZd/Op6biqnufUh3brFKJ4l9Ywdq7EyqN7a6bhnEkN1s6Kcnm1ZMSqV87/fqh099+5bQEdVY+eOIUF4YiieL8wEkj6WyiZdPQEmGpqZE2RHjhocZVmHHe1t1WffZsT1F5Yu6/MivD7Lw4K8EcYH1WWbi1SzSLEJxVWgXOYYmqPLrLVAIqaTOmUpL167Vcop6oep2bt9sT1RLUISo8cerYVctPdcI2qZRoSSk4ADWKZy9stIaNhPUQa8fbQrbVOvk2rnyVd6byyzVfBjSlalWEU6oVFQCrlJrSWLtEKQ/l1uVkkm5WlijY3D+gRH1Vvnt3v31RFpBOTcJx87G2Y6gWrW0TULys/dIb7FiN4LsDFOMWvYCTpUpz+1bARUHN7TsBB0Vp7tcIVBVgbR+wLGuHW1ZRDt4FLMd2b91uRjZGsIFsjET9bdvtycaGdG7eIiXoj2ayqRWtJhsjXtZ+IZtdyEb77q8lGz3mJ33w425UU8bYjWhKGHd70kxN6VXTTKmgnXut/2+aKY3t9jSpiEbgsI3aGiPb3NkYofrm5r7Xubrubd/f2MBOi6NPLb/VNjhVgrXtjRFeUX/hm7XPlwpfHe7p0paUs0y0/e92uEDW/Y6jZWYWLwpyql/mmEWPlESGL3I/mGWNTVxn+3Q6NTVZ2+yU6t3pdyqxDt7ylCxwux6JD9U2fF5tZJt7HyNU3/707zq9/ravuTjAThOkTy2/rW2FmhRqGyKjtALz0hZtTFOFz34aTRUW7N0clWF27o9KMAdokWqKtLZLKhGA0yhVYh28VypZ4NITjx836paWsmvoSQs10NPth09bd0sOsEtP6tTyuU5xaD1PbaJZT1i5tvNcaR3KeRPKOuWXG8HVtKd89lNpT1mwP+2VYHanPRvmELRXXfz1tGcTi0t7VViHpz3bAneVuDHvLWWbV4nNvAe1uxPvOcDOKtEwR/Ft7e3wJoXaVaPDT+uUX9qyytXjT+anZQLuv3p0gXZ+Wu4AHeB5+VYs5dKAs3j8C1jKtcCw1MrGCHdjS0TGbISbi3B/F5/o0lOGJBR3pJn5sGTC1+08qtjk5BXp7dRQFU1WZ3zhpNf9Yu/YcreYel2tuKugwIt31lbH0LGiEreLQYR+LH9+1i+gp5zHX5m+fe9/yV9sL21SyjXCbDGOUhoMNth70qrR2uQt9Drdde+E1umtfdurQbH55YsGxcYHNiZLSoFwSuUL+dcl+fvZ+dmZSZNqZ64dpsa8dTd6G+K37qbM5pP7Zc3k1vUZDfNbr2oRiYIw9DGP4IKoSjMVzxfFZfYrfW7nm7MSIrlI6ThYpZoQqGwug1ZxfVS7XAP/lbwgryRuZguWSMp+s/vcqn/c5Gmuh/JZhnp3bIh7vWngPyTXvd5t7wJK+SGxtpbJVAQA2CrUgBhS3AXneRCLwQB3rw0G5PKS+IMBznEwUGysp+v9DwcFMAM=
+\ No newline at end of file
+diff --git a/third_party/gpus/rocm_configure.bzl b/third_party/gpus/rocm_configure.bzl
+index 30760f4f..336c1e96 100644
+--- a/third_party/gpus/rocm_configure.bzl
++++ b/third_party/gpus/rocm_configure.bzl
+@@ -29,7 +29,6 @@ load(
+     "which",
+ )
+
+-
+ _GCC_HOST_COMPILER_PATH = "GCC_HOST_COMPILER_PATH"
+ _GCC_HOST_COMPILER_PREFIX = "GCC_HOST_COMPILER_PREFIX"
+ _ROCM_TOOLKIT_PATH = "ROCM_PATH"
+@@ -175,17 +174,23 @@ def _rocm_include_path(repository_ctx, rocm_config, bash_bin):
+     """
+     inc_dirs = []
+
++    hsa_path = get_host_environ(repository_ctx, "HSA_ROCR_DEV_PATH")
++    hip_path = get_host_environ(repository_ctx, "HIP_PATH")
++    rocprim_path = get_host_environ(repository_ctx, "ROCPRIM_PATH")
++    rocsolver_path = get_host_environ(repository_ctx, "ROCSOLVER_PATH")
++    rocblas_path = get_host_environ(repository_ctx, "ROCBLAS_PATH")
++    llvm_path = get_host_environ(repository_ctx, "LLVM_PATH")
+     # Add HSA headers (needs to match $HSA_PATH)
+-    inc_dirs.append(rocm_config.rocm_toolkit_path + "/hsa/include")
++    inc_dirs.append(hsa_path + "/include/hsa")
+
+     # Add HIP headers (needs to match $HIP_PATH)
+     inc_dirs.append(rocm_config.rocm_toolkit_path + "/hip/include")
+     if int(rocm_config.rocm_version_number) >= 50200:
+         inc_dirs.append(rocm_config.rocm_toolkit_path + "/include")
+-        inc_dirs.append(rocm_config.rocm_toolkit_path + "/include/hip")
+-        inc_dirs.append(rocm_config.rocm_toolkit_path + "/include/rocprim")
+-        inc_dirs.append(rocm_config.rocm_toolkit_path + "/include/rocsolver")
+-        inc_dirs.append(rocm_config.rocm_toolkit_path + "/include/rocblas")
++        inc_dirs.append(hip_path + "/include/hip")
++        inc_dirs.append(rocprim_path + "/include/rocprim")
++        inc_dirs.append(rocsolver_path + "/include/rocsolver")
++        inc_dirs.append(rocblas_path + "/include/rocblas")
+
+     # Add HIP-Clang headers (realpath relative to compiler binary)
+     rocm_toolkit_path = realpath(repository_ctx, rocm_config.rocm_toolkit_path, bash_bin)
+@@ -198,7 +203,7 @@ def _rocm_include_path(repository_ctx, rocm_config, bash_bin):
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/14.0.0/include")
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/15.0.0/include")
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/16.0.0/include")
+-    inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/17.0.0/include")
++    inc_dirs.append(llvm_path + "/lib/clang/17.0.0/include/")
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/17/include")
+     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/18/include")
+
+@@ -224,7 +229,8 @@ def _amdgpu_targets(repository_ctx, rocm_toolkit_path, bash_bin):
+     """Returns a list of strings representing AMDGPU targets."""
+     amdgpu_targets_str = get_host_environ(repository_ctx, _TF_ROCM_AMDGPU_TARGETS)
+     if not amdgpu_targets_str:
+-        cmd = "%s/bin/rocm_agent_enumerator" % rocm_toolkit_path
++        rocminfo_path = get_host_environ(repository_ctx, "ROCMINFO_PATH")
++        cmd = rocminfo_path + "/bin/rocm_agent_enumerator"
+         result = execute(repository_ctx, [bash_bin, "-c", cmd])
+         targets = [target for target in result.stdout.strip().split("\n") if target != "gfx000"]
+         targets = {x: None for x in targets}
+@@ -345,23 +351,33 @@ def _find_libs(repository_ctx, rocm_config, hipfft_or_rocfft, miopen_path, rccl_
+     Returns:
+       Map of library names to structs of filename and path
+     """
++    hip_path = get_host_environ(repository_ctx, "HIP_PATH")
++    rocblas_path = get_host_environ(repository_ctx, "ROCBLAS_PATH")
++    hipfft_path = get_host_environ(repository_ctx, "HIPFFT_PATH")
++    hiprand_path = get_host_environ(repository_ctx, "HIPRAND_PATH")
++    hipsparse_path = get_host_environ(repository_ctx, "HIPSPARSE_PATH")
++    roctracer_path = get_host_environ(repository_ctx, "ROCTRACER_PATH")
++    rocsolver_path = get_host_environ(repository_ctx, "ROCSOLVER_PATH")
++    hipsolver_path = get_host_environ(repository_ctx, "HIPSOLVER_PATH")
++    hipblas_path = get_host_environ(repository_ctx, "HIPBLAS_PATH")
++
+     libs_paths = [
+         (name, _rocm_lib_paths(repository_ctx, name, path))
+         for name, path in [
+-            ("amdhip64", rocm_config.rocm_toolkit_path),
+-            ("rocblas", rocm_config.rocm_toolkit_path),
+-            (hipfft_or_rocfft, rocm_config.rocm_toolkit_path),
+-            ("hiprand", rocm_config.rocm_toolkit_path),
+-            ("MIOpen", miopen_path),
++            ("amdhip64", hip_path + "/lib"),
++            ("rocblas", rocblas_path),
++            (hipfft_or_rocfft, hipfft_path),
++            ("hiprand", hiprand_path),
++           ("MIOpen", miopen_path),
+             ("rccl", rccl_path),
+-            ("hipsparse", rocm_config.rocm_toolkit_path),
+-            ("roctracer64", rocm_config.rocm_toolkit_path),
+-            ("rocsolver", rocm_config.rocm_toolkit_path),
++            ("hipsparse", hipsparse_path),
++            ("roctracer64", roctracer_path + "/lib"),
++            ("rocsolver", rocsolver_path),
+         ]
+     ]
+     if int(rocm_config.rocm_version_number) >= 40500:
+-        libs_paths.append(("hipsolver", _rocm_lib_paths(repository_ctx, "hipsolver", rocm_config.rocm_toolkit_path)))
+-        libs_paths.append(("hipblas", _rocm_lib_paths(repository_ctx, "hipblas", rocm_config.rocm_toolkit_path)))
++        libs_paths.append(("hipsolver", _rocm_lib_paths(repository_ctx, "hipsolver", hipsolver_path)))
++        libs_paths.append(("hipblas", _rocm_lib_paths(repository_ctx, "hipblas", hipblas_path)))
+
+     # hipblaslt may be absent even in versions of ROCm where it exists
+     # (it is not installed by default in some containers). Autodetect.
+@@ -578,12 +594,13 @@ def _create_local_rocm_repository(repository_ctx):
+     hipfft_or_rocfft = "rocfft" if rocm_version_number < 40100 else "hipfft"
+
+     # For ROCm 5.2 and above, find MIOpen and RCCL in the main rocm lib path
+-    miopen_path = rocm_config.rocm_toolkit_path + "/miopen" if rocm_version_number < 50200 else rocm_config.rocm_toolkit_path
+-    rccl_path = rocm_config.rocm_toolkit_path + "/rccl" if rocm_version_number < 50200 else rocm_config.rocm_toolkit_path
++    miopen_path = get_host_environ(repository_ctx, "MIOPEN_PATH") + "/miopen"
++    rccl_path = get_host_environ(repository_ctx, "RCCL_PATH") + "/rccl"
+
+     # Copy header and library files to execroot.
+     # rocm_toolkit_path
+     rocm_toolkit_path = rocm_config.rocm_toolkit_path
++    rocblas_path = get_host_environ(repository_ctx, "ROCBLAS_PATH")
+     copy_rules = [
+         make_copy_dir_rule(
+             repository_ctx,
+@@ -597,6 +614,53 @@ def _create_local_rocm_repository(repository_ctx):
+     # explicitly copy (into the local_config_rocm repo) the $ROCM_PATH/hiprand/include and
+     # $ROCM_PATH/rocrand/include dirs, only once the softlink to them in $ROCM_PATH/include
+     # dir has been removed. This removal will happen in a near-future ROCm release.
++    rocm_parts = [
++        "rocblas",
++       "rocprim",
++        "hipfft",
++        "hipsparse",
++        "rocsolver",
++        "hipsolver",
++        "hipblas",
++       "rccl",
++       "miopen",
++       "roctracer",
++    ]
++    rocm_includes = ""
++
++    for part in rocm_parts:
++        part_cap = part.upper().replace("-","_")
++        part_path = get_host_environ(repository_ctx, part_cap + "_PATH")
++        rocm_includes = rocm_includes + '":' + part + '-include",\n'
++        rocm_includes = rocm_includes + '":' + part + '-include-second",\n'
++        copy_rules.append(
++            make_copy_dir_rule(
++                repository_ctx,
++                name = part + "-include",
++                src_dir = part_path + "/include/" + part,
++                out_dir = "rocm/include",
++            ),
++        )
++        copy_rules.append(
++            make_copy_dir_rule(
++                repository_ctx,
++                name = part + "-include-second",
++                src_dir = part_path + "/include/" + part,
++                out_dir = "rocm/include/" + part,
++            ),
++        )
++    hipcub_path = get_host_environ(repository_ctx, "HIPCUB_PATH")
++    copy_rules.append(
++        make_copy_dir_rule(
++            repository_ctx,
++            name = "hipcub-include",
++            src_dir = hipcub_path + "/include/hipcub/",
++            out_dir = "rocm/include/hipcub/",
++        )
++    )
++    rocm_includes = rocm_includes + '":hipcub-include",\n'
++
++    hiprand_path = get_host_environ(repository_ctx, "HIPRAND_PATH")
+     hiprand_include = ""
+     hiprand_include_softlink = rocm_config.rocm_toolkit_path + "/include/hiprand"
+     softlink_exists = files_exist(repository_ctx, [hiprand_include_softlink], bash_bin)
+@@ -606,11 +670,12 @@ def _create_local_rocm_repository(repository_ctx):
+             make_copy_dir_rule(
+                 repository_ctx,
+                 name = "hiprand-include",
+-                src_dir = rocm_toolkit_path + "/hiprand/include",
++                src_dir = hiprand_path + "/include/hiprand",
+                 out_dir = "rocm/include/hiprand",
+             ),
+         )
+
++    rocrand_path = get_host_environ(repository_ctx, "ROCRAND_PATH")
+     rocrand_include = ""
+     rocrand_include_softlink = rocm_config.rocm_toolkit_path + "/include/rocrand"
+     softlink_exists = files_exist(repository_ctx, [rocrand_include_softlink], bash_bin)
+@@ -620,11 +685,10 @@ def _create_local_rocm_repository(repository_ctx):
+             make_copy_dir_rule(
+                 repository_ctx,
+                 name = "rocrand-include",
+-                src_dir = rocm_toolkit_path + "/rocrand/include",
++                src_dir = rocrand_path + "/include/rocrand",
+                 out_dir = "rocm/include/rocrand",
+             ),
+         )
+-
+     rocm_libs = _find_libs(repository_ctx, rocm_config, hipfft_or_rocfft, miopen_path, rccl_path, bash_bin)
+     rocm_lib_srcs = []
+     rocm_lib_outs = []
+@@ -639,7 +703,8 @@ def _create_local_rocm_repository(repository_ctx):
+         outs = rocm_lib_outs,
+     ))
+
+-    clang_offload_bundler_path = rocm_toolkit_path + "/llvm/bin/clang-offload-bundler"
++    llvm_amdgpu_path = get_host_environ(repository_ctx, "LLVM_PATH")
++    clang_offload_bundler_path = llvm_amdgpu_path + "/bin/clang-offload-bundler"
+
+     # copy files mentioned in third_party/gpus/rocm/BUILD
+     copy_rules.append(make_copy_files_rule(
+@@ -683,9 +748,7 @@ def _create_local_rocm_repository(repository_ctx):
+         "%{roctracer_lib}": rocm_libs["roctracer64"].file_name,
+         "%{rocsolver_lib}": rocm_libs["rocsolver"].file_name,
+         "%{copy_rules}": "\n".join(copy_rules),
+-        "%{rocm_headers}": ('":rocm-include",\n' +
+-                            hiprand_include +
+-                            rocrand_include),
++        "%{rocm_headers}": ('":rocm-include",\n' + rocrand_include + hiprand_include + rocm_includes),
+     }
+     if rocm_libs["hipblaslt"] != None:
+         repository_dict["%{hipblaslt_lib}"] = rocm_libs["hipblaslt"].file_name
+@@ -749,16 +812,18 @@ def _create_local_rocm_repository(repository_ctx):
+         tpl_paths["crosstool:hipcc_cc_toolchain_config.bzl"],
+     )
+
++    hip_path = get_host_environ(repository_ctx, "HIP_PATH")
++    hsa_rocr_dev_path = get_host_environ(repository_ctx, "HSA_ROCR_DEV_PATH")
+     repository_ctx.template(
+         "crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc",
+         tpl_paths["crosstool:clang/bin/crosstool_wrapper_driver_rocm"],
+         {
+             "%{cpu_compiler}": str(cc),
+-            "%{hipcc_path}": rocm_config.rocm_toolkit_path + "/bin/hipcc",
++            "%{hipcc_path}": hip_path + "/bin/hipcc",
+             "%{hipcc_env}": _hipcc_env(repository_ctx),
+-            "%{rocr_runtime_path}": rocm_config.rocm_toolkit_path + "/lib",
++            "%{rocr_runtime_path}": hsa_rocr_dev_path + "/lib",
+             "%{rocr_runtime_library}": "hsa-runtime64",
+-            "%{hip_runtime_path}": rocm_config.rocm_toolkit_path + "/lib",
++            "%{hip_runtime_path}": hip_path + "/lib",
+             "%{hip_runtime_library}": "amdhip64",
+             "%{crosstool_verbose}": _crosstool_verbose(repository_ctx),
+             "%{gcc_host_compiler_path}": str(cc),
diff --git a/var/spack/repos/builtin/packages/py-tensorflow/package.py b/var/spack/repos/builtin/packages/py-tensorflow/package.py
index 6668507186..1a3c3154a0 100644
--- a/var/spack/repos/builtin/packages/py-tensorflow/package.py
+++ b/var/spack/repos/builtin/packages/py-tensorflow/package.py
@@ -8,12 +8,36 @@ import tempfile
 
 from spack.package import *
 
+rocm_dependencies = [
+    "hip",
+    "rocrand",
+    "rocblas",
+    "rocfft",
+    "hipfft",
+    "rccl",
+    "hipsparse",
+    "rocprim",
+    "llvm-amdgpu",
+    "hsa-rocr-dev",
+    "rocminfo",
+    "hipsolver",
+    "hiprand",
+    "rocsolver",
+    "hipsolver",
+    "hipblas",
+    "hipcub",
+    "rocm-core",
+    "roctracer-dev",
+    "miopen-hip",
+]
+
 
 class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     """TensorFlow is an open source machine learning framework for everyone."""
 
     homepage = "https://www.tensorflow.org"
     url = "https://github.com/tensorflow/tensorflow/archive/v2.3.1.tar.gz"
+    git = "https://github.com/tensorflow/tensorflow.git"
 
     maintainers("adamjstewart", "aweits")
     import_modules = ["tensorflow"]
@@ -21,9 +45,19 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     license("Apache-2.0")
 
     version("2.16.1", sha256="c729e56efc945c6df08efe5c9f5b8b89329c7c91b8f40ad2bb3e13900bd4876d")
+    version(
+        "2.16-rocm-enhanced",
+        git="https://github.com/ROCm/tensorflow-upstream.git",
+        branch="r2.16-rocm-enhanced",
+    )
     version("2.15.1", sha256="f36416d831f06fe866e149c7cd752da410a11178b01ff5620e9f265511ed57cf")
     version("2.15.0", sha256="9cec5acb0ecf2d47b16891f8bc5bc6fbfdffe1700bdadc0d9ebe27ea34f0c220")
     version("2.14.1", sha256="6b31ed347ed7a03c45b906aa41628ac91c3db7c84cb816971400d470e58ba494")
+    version(
+        "2.14-rocm-enhanced",
+        git="https://github.com/ROCm/tensorflow-upstream.git",
+        branch="r2.14-rocm-enhanced-nohipblaslt-build",
+    )
     version("2.14.0", sha256="ce357fd0728f0d1b0831d1653f475591662ec5bca736a94ff789e6b1944df19f")
     version("2.13.1", sha256="89c07aebd4f41fbe0d08cc88aef00305542134f2f16d3b62918dc3c1182f33e2")
     version("2.13.0", sha256="e58c939079588623e6fa1d054aec2f90f95018266e0a970fd353a5244f5173dc")
@@ -286,20 +320,9 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     # depends_on('android-sdk', when='+android')
 
     with when("+rocm"):
-        depends_on("hip")
-        depends_on("rocrand")
-        depends_on("rocblas")
-        depends_on("rocfft")
-        depends_on("hipfft")
-        depends_on("rccl", when="+nccl")
-        depends_on("hipsparse")
-        depends_on("hipcub")
-        depends_on("rocsolver")
-        depends_on("rocprim")
-        depends_on("miopen-hip")
-        depends_on("llvm-amdgpu")
-        depends_on("hsa-rocr-dev")
-        depends_on("rocminfo")
+        for pkg_dep in rocm_dependencies:
+            depends_on(f"{pkg_dep}@6.0:", when="@2.14:")
+            depends_on(pkg_dep)
 
     # Check configure and configure.py to see when these variants are supported
     conflicts("+mkl", when="platform=darwin", msg="Darwin is not yet supported")
@@ -354,9 +377,11 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     conflicts("platform=darwin target=aarch64:", when="@:2.4")
     # https://github.com/tensorflow/tensorflow/pull/39225
     conflicts("target=aarch64:", when="@:2.2")
-    conflicts("~rocm", when="@2.7.4-rocm-enhanced,2.11.0-rocm-enhanced")
-    conflicts("+rocm", when="@:2.7.4-a,2.7.4.0:2.11.0-a,2.11.0.0:")
-
+    conflicts(
+        "~rocm",
+        when="@2.7.4-rocm-enhanced,2.11.0-rocm-enhanced,2.14-rocm-enhanced,2.16-rocm-enhanced",
+    )
+    conflicts("+rocm", when="@:2.7.4-a,2.7.4.0:2.11.0-a,2.11.0.0:2.14-a,2.14-z:2.16-a,2.16-z:")
     # wheel 0.40 upgrades vendored packaging, trips over tensorflow-io-gcs-filesystem identifier
     conflicts("^py-wheel@0.40:", when="@2.11:2.13")
 
@@ -403,6 +428,10 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     # and https://github.com/abseil/abseil-cpp/issues/1665
     patch("absl_neon.patch", when="@2.16.1: target=aarch64:")
 
+    patch("Find_ROCm_Components_Individiually.2.16.patch", when="@2.16-rocm-enhanced +rocm")
+    patch("Find_ROCm_Components_Individiually.patch", when="@2.14-rocm-enhanced +rocm")
+    patch("set_jit_trueLT_false.patch", when="@2.14-rocm-enhanced: +rocm")
+
     phases = ["configure", "build", "install"]
 
     # https://www.tensorflow.org/install/source
@@ -512,6 +541,13 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
         # Do you wish to build TensorFlow with ROCm support?
         if "+rocm" in spec:
             env.set("TF_NEED_ROCM", "1")
+            env.set("TF_HIPBLASLT", "0")
+            env.set("MIOPEN_PATH", spec["miopen-hip"].prefix)
+            env.set("ROCTRACER_PATH", spec["roctracer-dev"].prefix)
+            env.set("LLVM_PATH", spec["llvm-amdgpu"].prefix)
+            for pkg_dep in rocm_dependencies:
+                pkg_dep_cap = pkg_dep.upper().replace("-", "_")
+                env.set(f"{pkg_dep_cap}_PATH", spec[pkg_dep].prefix)
         else:
             env.set("TF_NEED_ROCM", "0")
 
@@ -704,6 +740,11 @@ class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
             with open(".tf_configure.bazelrc", mode="a") as f:
                 f.write('build --action_env LD_LIBRARY_PATH="' + slibs + '"')
 
+        if spec.satisfies("@2.16-rocm-enhanced +rocm"):
+            filter_file(
+                "/usr/lib/llvm-17/bin/clang", spec["llvm-amdgpu"].prefix.bin.clang, ".bazelrc"
+            )
+
         filter_file("build:opt --copt=-march=native", "", ".tf_configure.bazelrc")
         filter_file("build:opt --host_copt=-march=native", "", ".tf_configure.bazelrc")
 
diff --git a/var/spack/repos/builtin/packages/py-tensorflow/set_jit_trueLT_false.patch b/var/spack/repos/builtin/packages/py-tensorflow/set_jit_trueLT_false.patch
new file mode 100644
index 0000000000..ee58be0745
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-tensorflow/set_jit_trueLT_false.patch
@@ -0,0 +1,22 @@
+diff --git a/tensorflow/core/kernels/mlir_generated/build_defs.bzl b/tensorflow/core/kernels/mlir_generated/build_defs.bzl
+index 1cc3d335577..327e6cda821 100644
+--- a/tensorflow/core/kernels/mlir_generated/build_defs.bzl
++++ b/tensorflow/core/kernels/mlir_generated/build_defs.bzl
+@@ -355,7 +355,7 @@ def _gen_kernel_library(
+                 extra_args = extra_args,
+                 host_triple = host_triple,
+                 gpu_archs = gpu_archs,
+-                jit = jit,
++                jit = True,
+                 max_supported_rank = max_supported_rank,
+                 mlir_op = "{op}_{name}_{platform}_{type}_{output_type}.mlir".format(
+                     op = op,
+@@ -366,7 +366,7 @@ def _gen_kernel_library(
+                 ),
+                 tile_size = typed_tile_size,
+                 unroll_factors = typed_unroll_factors,
+-                jit_i64_indexed_for_large_tensors = jit_i64_indexed_for_large_tensors,
++                jit_i64_indexed_for_large_tensors = False,
+             )
+ 
+             # We have to use a sh_test instead of build_test because it doesn't properly find the dependent targets.