Adding mathlib include paths

Author: renjithravindrankannath

caffe2/CMakeLists.txt

@@ -1373,13 +1373,6 @@ if(USE_ROCM)
     set(ROCM_SOURCE_DIR "/opt/rocm")
   endif()
   message(INFO "caffe2 ROCM_SOURCE_DIR = ${ROCM_SOURCE_DIR}")
-  target_include_directories(torch_hip PRIVATE
-    ${ROCM_SOURCE_DIR}/include
-    ${ROCM_SOURCE_DIR}/hcc/include
-    ${ROCM_SOURCE_DIR}/rocblas/include
-    ${ROCM_SOURCE_DIR}/hipsparse/include
-    ${ROCM_SOURCE_DIR}/include/rccl/
-    )
   if(USE_FLASH_ATTENTION)
     target_compile_definitions(torch_hip PRIVATE
         USE_FLASH_ATTENTION
@@ -1713,7 +1706,7 @@ if(USE_ROCM)
   target_link_libraries(torch_hip PRIVATE ${Caffe2_HIP_DEPENDENCY_LIBS})
 
   # Since PyTorch files contain HIP headers, this is also needed to capture the includes.
-  target_include_directories(torch_hip PRIVATE ${Caffe2_HIP_INCLUDE})
+  target_include_directories(torch_hip PRIVATE ${Caffe2_HIP_INCLUDE} ${ROCM_INCLUDE})
   target_include_directories(torch_hip INTERFACE $<INSTALL_INTERFACE:include>)
 endif()
 

cmake/public/LoadHIP.cmake

@@ -165,6 +165,12 @@ if(HIP_FOUND)
   endif()
   find_package_and_print_version(hipblaslt REQUIRED)
 
+  list(APPEND ROCM_INCLUDE ${rocthrust_INCLUDE_DIR})
+  list(APPEND ROCM_INCLUDE ${rocprim_INCLUDE_DIR})
+  list(APPEND ROCM_INCLUDE ${hipcub_INCLUDE_DIR})
+  list(APPEND ROCM_INCLUDE ${rocRAND_INCLUDE_DIR})
+  list(APPEND ROCM_INCLUDE ${INTERFACE_INCLUDE_DIRECTORIES})
+
   if(UNIX)
     find_package_and_print_version(rccl)
     find_package_and_print_version(hsa-runtime64 REQUIRED)