Merge pull request tensorflow#1668 from ROCmSoftwarePlatform/develop-…

…upstream-52-hip-include

[develop-upstream] Merge in ROCm5.2 HIP include file changes from staging

tensorflow/compiler/xla/service/gpu/nccl_utils.h

@@ -27,8 +27,12 @@ limitations under the License.
 
 // Common place for all collective thunks to include nccl/rccl headers.
 #if TENSORFLOW_USE_ROCM
+#if (TF_ROCM_VERSION >= 50200)
 #include "rocm/include/rccl/rccl.h"
 #else
+#include "rocm/include/rccl.h"
+#endif
+#else
 #include "third_party/nccl/nccl.h"
 #endif
 

tensorflow/core/kernels/nccl_ops.cc

@@ -20,7 +20,11 @@ limitations under the License.
 #if GOOGLE_CUDA
 #include "third_party/nccl/nccl.h"
 #elif TENSORFLOW_USE_ROCM
+#if (TF_ROCM_VERSION >= 50200)
 #include "rocm/include/rccl/rccl.h"
+#else
+#include "rocm/include/rccl.h"
+#endif
 #endif
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/nccl/nccl_manager.h"

tensorflow/core/nccl/nccl_manager.h

@@ -30,7 +30,11 @@ limitations under the License.
 #if GOOGLE_CUDA
 #include "third_party/nccl/nccl.h"
 #elif TENSORFLOW_USE_ROCM
+#if (TF_ROCM_VERSION >= 50200)
 #include "rocm/include/rccl/rccl.h"
+#else
+#include "rocm/include/rccl.h"
+#endif
 #include "tensorflow/core/common_runtime/gpu_device_context.h"
 #endif
 #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h"

third_party/gpus/find_rocm_config.py

@@ -135,8 +135,17 @@ def hipruntime_version_number(path):
 def _find_miopen_config(rocm_install_path):
 
   def miopen_version_numbers(path):
-    version_file = os.path.join(path, "miopen/include/miopen/version.h")
-    if not os.path.exists(version_file):
+    possible_version_files = [
+            "include/miopen/version.h",         # ROCm 5.2 and prior
+            "miopen/include/miopen/version.h",  # ROCm 5.1 and prior
+    ]
+    version_file = None
+    for f in possible_version_files:
+      version_file_path = os.path.join(path, f)
+      if os.path.exists(version_file_path):
+        version_file = version_file_path
+        break
+    if not version_file:
       raise ConfigError(
           'MIOpen version file "{}" not found'.format(version_file))
     major = _get_header_version(version_file, "MIOPEN_VERSION_MAJOR")

third_party/gpus/find_rocm_config.py.gz.base64

@@ -1 +1 @@
-eJztW21v2zgS/q5fQSgoKm8cJSlugUMOOcCbZlHftUlgZ7tYNIVB27TNrSz6SCppUPS/7wxJyZIs2YqjvX5JgDa2NPNwOC/PjBTpgFyI1aPk84Umb07enJDbBSO3LFZC/hqJB9JL9EJIFZJeFJEBiikyYIrJezYNvQPvgLznExBnU5LEUyaJBv3eik7glzvTJR+ZVFzE5E14QgIU8N0pv/MvQHgUCVnSRxILTRLFAIIrMuMRI+zrhK004TGZiOUq4jSeMPLA9cIs40DADPKHgxBjTUGagvwKvs3ycoRqYzD+LLRenR0fPzw8hNQYGwo5P46soDp+37+4vBpeHoHBRuW3OGJKEcn+l3AJWx0/EroCeyZ0DFZG9IEISehcMjinBdr7ILnm8bxLlJjpByoZoEy50pKPE11wVmod7DkvAO6iMfF7Q9If+uSX3rA/7ALG7/3bd9e/3ZLfe4NB7+q2fzkk1wNycX31tn/bv76Cb7+S3tUf5L/9q7ddwsBVsAz7upJoPxjJ0Y0mdGTIWMGAmbAGqRWb8BmfwL7ieULnjMzFPZMxbIesmFxyhcFUYN4UUCK+5Jpqc2RjU7jMeas/nu/7N5LHmIbXF0tYfiypfERjyIJRXH8KIZpoITkzNpJ7m32QUgIMRMeaXT4qzZah52HCq4nkkGeKUQm5oIwr6uAxMVURpQsRR69p5cHBJabAlGl0VWxczGVqhAFaWftRfyLiGZ8n0jgQ9ZSeikSHxqoVxUQXKThmiIsNptlCimS+wCRh8T2XIl6yWJN7KrlJygDs/zC66d2+64RefwbFBeciPi0tyZ1bunY71g+pgcYcJqUJtWQ6kSbsBA6BgyZiyor+0/QLs/tKY/CYsxiKBk9ldlXaHebxIiG+2GBY39t4pjGxgTDVvqByeoT2TCGGGureU8k4nwczKZZkTJVzqiOGtW2ZvSEBX61NBPcAK3mZoHETlOWxWOljKSZLH0USpD8KtmiI+4wmEe4nSpiH2ep5UHNCQvhE+kmo9BPwgvsEmeR53iSiUKcXJkSX6OXg0lAghKpz5hGwXqEYrEJGc6ZHbrkRmjLCrQVGzMYqb2ZeyQhDUmkaRTklsPVtmrXW02nIXdiWxCnZ1EHNEDcI65URyXmtfSDOZ8TPfOxjCIUKXTagJdWAa5lPOe3PIH9QJ29MlIxGZukNoc7aVRvn8g7D5iMU12zkyngUJ8sxk8GS/ilkl4DH8BeoTRZ5/5+ewA/5iRgxcojf8RtKwzcjnl/GJnm6RoBWdElMlywNz8DFA9h8BQYDpQMQqJuAgRcv8hzlQoPVEyEPwXkuQrFiKbIvfegQMZQNMPu5n+jZ0T/9jvX/Em0DH0oWmo+B9A/sQuSVIofB3fSw45NXxrquwe8YPYiskbcoxJYB4JiD4RwoaxWcdtxJ5yXgmsDIdTx0HeVQpfkKeJ2u7H/7bsrNdrW7+HUImwPkwPiIHDrY4o9PzFyBDRB5BrxgJotv36Fv3sV+CmHSoR6CAx3avVhGQBYEu9mcoXehaGjkd9JYgrFTm/GWaysS7wy3isLmVDGrVOBEcG3IO8WBFzMZy3rn5FNmqs/jSZRMmSl0YEFpP6Ua4QLCDCViCvjn8E21Xo30adqwhDR6n83/eVvAlCsRM3Mcc22GLq62OkuK3MFysf4peJqfszRPIKnS0+wrjEgq2EDonGXbKhm3IZoJjoEYvqRZizHNi6Z4m+mYSxHfeClt7WZBxLFzBrgB8jXNr1JiVTuoYyqAOL44r2SFvELX0ejHy8EQhr/Rh95/rge+dZtlmT0w+lcZxsqxwFMxgJgv3jkMV+QVTIl7rTiMrFNVFBXsnXYeW2ag+M252a9A8M/2IHPA++6VuoRdrVDsC76SSaz5kjUo+ZxwyYI96x4A8d/ugkfBZkovdb+17t/1b8jAxvBvKP8DuAiHNpPr5TDg2sH7K150wpiZtkT0PBiTVl6XPCwE2G2nT4sFSEFOxJIEjCE4jBySwhks/Y5Tg1lDMn5vLxHTAcaSih0NjyY0miQwC4J1TAIHKG6uzmH4F2ZotjgPFLou1B22S2HAxonCgUXhFL6isH+4roTpH9IBr1PVUxlwY29PJ8ANJxS567QwxBlcDNRG0ecJqLbIgYbqCaByQs0T0MaaBRpachzvGlCQE9wyd5SqqaJMfQuSEYr7uuaTTr7GthRyp1HNvf7Qvwb8YrmZiTCruWweLKB3nppQsNDN5dVzc6qMsl9bLaE8t7HWxL26tRayKZ/blSj+WS5Yz260hbXLc/UYLpKbjdZGsr3pGuGOceyXMY3SA0c72+5a0yI1RHjpwVt7sHNeu/33iVQBkfrlfW/YwgBehNl7Bi/AtDCGVxZQ7SSeK8zSMF6B0y5hFFcvM4aESmrGGEaynbHcoR2XUMtVXuAJo1ACqFE8eaGHHfQgc3f/t9PDVj4oZkPjUhz0rt6mpVgswiKgq51cipZqpyIlYYisydWdE2RxqXKhzGa6WZ2gYHuNFdDcr0bdFMWL2i8tdM8aAd/90A5aSiWj+fQGWkZBzaf3zxKK0Xxe+6yqktruua69EgFsorTeO9drl+9rNWMEJ9gSI1g092s3IzjxovYLI+zDCNZ3P5QRrAkf7cEP+/FBEWM/Nihg3DyTC2rqo5oLClVXuqn0d3NBYe3ydKAlnYDDGg0ITra9GcECrj9tnRCc8Ib6Cxs8cT6wbvvRF9m3g97F5aCFy+wyUO5C+2CdtdmTRILZv/Yu6D0j1FVz6gs3vudp5eQ5w0JNxdTOC4VqLI0MlVitTw0FC8qDg8J7+w3/HuZk2xsfLOD6U6MhwiltwLyMEnuOEtZ9P3qasFY8e6Aowuw9UxRgWhgrakqndrIolGVpuKjEan2+KFhQJg0R3TeaMNayLZKGAVzfns8ONWOPVLt0k38Xyuk2lF3KLzcBNzho7cUnE0+bpGMseD7pFGD2J508TBukU1169aSTL+sy6VRhtU86eQvK1zWNWWct2951jav47FOjO6Almtil/DKq7LrO2ZMx2mQNiNXw+v3HVq5zykB7/0GxBNTCnxSfxB3l4ixd5vwfuKNsgeOOjWd8y89oUzLlE3xOHl9VEDNbicaQmOHLFW4/PJ6JXQ/RVzyv79U9WFNFXjsK4PUwew1i/eKDebImu/414NnjNZtr2FgxhW8+QIS+r79+so9iaiGiL1wbaf9z+ohn4Xn7VCNMVlOqWTBq8Bh1p0aryfOYdbq7HqKq09v5YMoWxe1/n96iuPXuvMuSQiBKpfKZ/Puc/OPk9OTEpUm1M3cuU2PerhuGW+K36/ZB8839vGNzu7r/lv3tVs0RiYFw9LGk0B9NaWr5eJZ13S/ssZu+ZxATJaRm02CTakKgsqUKOlm7NC9sBf4rdUZeKXwvI1gjGfvdi5S5+sf3lVx7VI8qtC96hfjaIgv8u/hyMLgenEEp38W5tySUlgEAdjI1IAaNL3R4HsRiNMIXMUYjcn5O/NEI9zgaGTa22/X+An4mbL4=
+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==

third_party/gpus/rocm_configure.bzl

@@ -179,7 +179,10 @@ def _rocm_include_path(repository_ctx, rocm_config, bash_bin):
 
     # Add HIP headers (needs to match $HIP_PATH)
     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")
     else:
         inc_dirs.append(rocm_config.rocm_toolkit_path + "/hip/include")
 
@@ -317,7 +320,7 @@ def _select_rocm_lib_paths(repository_ctx, libs_paths, bash_bin):
 
     return libs
 
-def _find_libs(repository_ctx, rocm_config, hipfft_or_rocfft, bash_bin):
+def _find_libs(repository_ctx, rocm_config, hipfft_or_rocfft, miopen_path, rccl_path, bash_bin):
     """Returns the ROCm libraries on the system.
 
     Args:
@@ -335,8 +338,8 @@ def _find_libs(repository_ctx, rocm_config, hipfft_or_rocfft, bash_bin):
             ("rocblas", rocm_config.rocm_toolkit_path),
             (hipfft_or_rocfft, rocm_config.rocm_toolkit_path),
             ("hiprand", rocm_config.rocm_toolkit_path),
-            ("MIOpen", rocm_config.rocm_toolkit_path + "/miopen"),
-            ("rccl", rocm_config.rocm_toolkit_path + "/rccl"),
+            ("MIOpen", miopen_path),
+            ("rccl", rccl_path),
             ("hipsparse", rocm_config.rocm_toolkit_path),
             ("roctracer64", rocm_config.rocm_toolkit_path + "/roctracer"),
             ("rocsolver", rocm_config.rocm_toolkit_path),
@@ -552,7 +555,11 @@ def _create_local_rocm_repository(repository_ctx):
     # For ROCm 4.1 and above use hipfft, older ROCm versions use rocfft
     rocm_version_number = int(rocm_config.rocm_version_number)
     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
+
     # Copy header and library files to execroot.
     # rocm_toolkit_path
     rocm_toolkit_path = rocm_config.rocm_toolkit_path
@@ -570,18 +577,6 @@ def _create_local_rocm_repository(repository_ctx):
             src_dir = rocm_toolkit_path + "/rocblas/lib/library",
             out_dir = "rocm/lib/rocblas/lib/library",
         ),
-        make_copy_dir_rule(
-            repository_ctx,
-            name = "miopen-include",
-            src_dir = rocm_toolkit_path + "/miopen/include",
-            out_dir = "rocm/include/miopen",
-        ),
-        make_copy_dir_rule(
-            repository_ctx,
-            name = "rccl-include",
-            src_dir = rocm_toolkit_path + "/rccl/include",
-            out_dir = "rocm/include/rccl",
-        )
     ]
 
     # explicitly copy (into the local_config_rocm repo) the $ROCM_PATH/hiprand/include and
@@ -615,7 +610,7 @@ def _create_local_rocm_repository(repository_ctx):
             ),
         )
 
-    rocm_libs = _find_libs(repository_ctx, rocm_config, hipfft_or_rocfft, bash_bin)
+    rocm_libs = _find_libs(repository_ctx, rocm_config, hipfft_or_rocfft, miopen_path, rccl_path, bash_bin)
     rocm_lib_srcs = []
     rocm_lib_outs = []
     for lib in rocm_libs.values():
@@ -671,8 +666,6 @@ def _create_local_rocm_repository(repository_ctx):
         "%{rocsolver_lib}": rocm_libs["rocsolver"].file_name,
         "%{copy_rules}": "\n".join(copy_rules),
         "%{rocm_headers}": ('":rocm-include",\n' +
-                            '":miopen-include",\n' +
-                            '":rccl-include",\n' +
                             hiprand_include +
                             rocrand_include),
     }