[HIPIFY][fix] CUDA and cuBLAS main headers correct handling

hipamd/hipify-clang/src/CUDA2HipMap.cpp

@@ -379,12 +379,12 @@ const std::map <llvm::StringRef, hipCounter> CUDA_INCLUDE_MAP{
     {"vector_types.h",       {"hip/hip_vector_types.h",   CONV_INCLUDE,             API_RUNTIME}},
 
     // CUBLAS includes
-    {"cublas.h",    {"hipblas.h", CONV_INCLUDE, API_BLAS}},
-    {"cublas_v2.h", {"hipblas.h", CONV_INCLUDE, API_BLAS}},
+    {"cublas.h",             {"hipblas.h",                CONV_INCLUDE_CUDA_MAIN_H, API_BLAS}},
+    {"cublas_v2.h",          {"hipblas.h",                CONV_INCLUDE_CUDA_MAIN_H, API_BLAS}},
 
     // CURAND includes
-    {"curand.h",             {"hiprand.h",        CONV_INCLUDE, API_RAND}},
-    {"curand_kernel.h",      {"hiprand_kernel.h", CONV_INCLUDE, API_RAND}},
+    {"curand.h",             {"hiprand.h",                CONV_INCLUDE,             API_RAND}},
+    {"curand_kernel.h",      {"hiprand_kernel.h",         CONV_INCLUDE,             API_RAND}},
 
     // HIP includes
     // TODO: uncomment this when hip/cudacommon.h will be renamed to hip/hipcommon.h

hipamd/hipify-clang/src/HipifyAction.cpp

@@ -149,23 +149,38 @@ void HipifyAction::InclusionDirective(clang::SourceLocation hash_loc,
     if (!SM.isWrittenInMainFile(hash_loc)) {
         return;
     }
+    if (!firstHeader) {
+        firstHeader = true;
+        firstHeaderLoc = hash_loc;
+    }
 
     const auto found = CUDA_INCLUDE_MAP.find(file_name);
     if (found == CUDA_INCLUDE_MAP.end()) {
-        if (!firstNotMainHeader) {
-            firstNotMainHeader = true;
-            firstNotMainHeaderLoc = hash_loc;
-        }
         return;
     }
 
     // Special-casing to avoid duplication of the hip_runtime include.
     bool secondMainInclude = false;
-    if (found->second.hipName == "hip/hip_runtime.h") {
-        if (insertedRuntimeHeader) {
-            secondMainInclude = true;
+    if (found->second.countType == CONV_INCLUDE_CUDA_MAIN_H) {
+        switch (found->second.countApiType) {
+            case API_DRIVER:
+            case API_RUNTIME:
+                if (insertedRuntimeHeader) {
+                    secondMainInclude = true;
+                    break;
+                }
+                insertedRuntimeHeader = true;
+                break;
+            case API_BLAS:
+                if (insertedBLASHeader) {
+                    secondMainInclude = true;
+                    break;
+                }
+                insertedBLASHeader = true;
+                break;
+            default:
+                break;
         }
-        insertedRuntimeHeader = true;
     }
 
     Statistics::current().incrementCounter(found->second, file_name.str());
@@ -356,8 +371,8 @@ void HipifyAction::EndSourceFileAction() {
         clang::SourceLocation sl;
         if (pragmaOnce) {
             sl = pragmaOnceLoc;
-        } else if (firstNotMainHeader) {
-            sl = firstNotMainHeaderLoc;
+        } else if (firstHeader) {
+            sl = firstHeaderLoc;
         } else {
             sl = SM.getLocForStartOfFile(SM.getMainFileID());
         }

hipamd/hipify-clang/src/HipifyAction.h

@@ -23,9 +23,10 @@ private:
     // not, we insert it at the top of the file when we finish processing it.
     // This approach means we do the best it's possible to do w.r.t preserving the user's include order.
     bool insertedRuntimeHeader = false;
-    bool firstNotMainHeader = false;
+    bool insertedBLASHeader = false;
+    bool firstHeader = false;
     bool pragmaOnce = false;
-    clang::SourceLocation firstNotMainHeaderLoc;
+    clang::SourceLocation firstHeaderLoc;
     clang::SourceLocation pragmaOnceLoc;
 
     /**

hipamd/tests/hipify-clang/headers_test_01.cu

@@ -1,6 +1,8 @@
 // RUN: %run_test hipify "%s" "%t" %cuda_args
 
 // CHECK: #include <hip/hip_runtime.h>
+// CHECK-NOT: #include <cuda_runtime.h>
+// CHECK: #include <stdio.h>
 #include <cuda.h>
-// CHECK-NOT: #include<cuda_runtime.h>
 #include <cuda_runtime.h>
+#include <stdio.h>

hipamd/tests/hipify-clang/headers_test_02.cu

@@ -1,6 +1,8 @@
 // RUN: %run_test hipify "%s" "%t" %cuda_args
 
-// CHECK: #include <hip/hip_runtime.h>
-#include <cuda_runtime.h>
-// CHECK-NOT: #include<cuda.h>
-#include <cuda.h>
+// CHECK: #include "hip/hip_runtime.h"
+// CHECK-NOT: #include "cuda_runtime.h"
+// CHECK: #include <stdio.h>
+#include "cuda.h"
+#include "cuda_runtime.h"
+#include <stdio.h>

hipamd/tests/hipify-clang/headers_test_06.cu

@@ -0,0 +1,8 @@
+// RUN: %run_test hipify "%s" "%t" %cuda_args
+
+// CHECK: #include <hipblas.h>
+// CHECK-NOT: #include <cublas_v2.h>
+// CHECK: #include <stdio.h>
+#include <cublas.h>
+#include <cublas_v2.h>
+#include <stdio.h>

hipamd/tests/hipify-clang/headers_test_07.cu

@@ -0,0 +1,8 @@
+// RUN: %run_test hipify "%s" "%t" %cuda_args
+
+// CHECK: #include "hipblas.h"
+// CHECK-NOT: #include "cublas.h"
+// CHECK: #include <stdio.h>
+#include "cublas_v2.h"
+#include "cublas.h"
+#include <stdio.h>

hipamd/tests/hipify-clang/headers_test_08.cu

@@ -0,0 +1,14 @@
+// RUN: %run_test hipify "%s" "%t" %cuda_args
+
+// CHECK: #include <hip/hip_runtime.h>
+// CHECK-NOT: #include <cuda_runtime.h>
+// CHECK: #include <iostream>
+// CHECK: #include "hipblas.h"
+// CHECK-NOT: #include "cublas.h"
+// CHECK: #include <stdio.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <iostream>
+#include "cublas_v2.h"
+#include "cublas.h"
+#include <stdio.h>