Merge pull request #1547 from emankov/clang

[HIPIFY][CUB][#1460] Implement cubFunctionTemplateDecl matcher [ROCm/clr commit: b4fd41169e]
2019-10-16 13:09:49 +03:00
@@ -61,6 +61,7 @@ const StringRef sCudaLaunchKernel = "cudaLaunchKernel";
 const StringRef sCudaHostFuncCall = "cudaHostFuncCall";
 const StringRef sCudaDeviceFuncCall = "cudaDeviceFuncCall";
 const StringRef sCubNamespacePrefix = "cubNamespacePrefix";
+const StringRef sCubFunctionTemplateDecl = "cubFunctionTemplateDecl";

 std::set<std::string> DeviceSymbolFunctions0 {
  {sCudaMemcpyToSymbol},
@@ -449,6 +450,41 @@ bool HipifyAction::cubNamespacePrefix(const mat::MatchFinder::MatchResult &Resul
  return false;
 }

+bool HipifyAction::cubFunctionTemplateDecl(const mat::MatchFinder::MatchResult &Result) {
+  if (auto *decl = Result.Nodes.getNodeAs<clang::FunctionTemplateDecl>(sCubFunctionTemplateDecl)) {
+    auto *Tparams = decl->getTemplateParameters();
+    bool ret = false;
+    for (size_t I = 0; I < Tparams->size(); ++I) {
+      const clang::ValueDecl *valueDecl = dyn_cast<clang::ValueDecl>(Tparams->getParam(I));
+      if (!valueDecl) continue;
+      clang::QualType QT = valueDecl->getType();
+      auto *t = QT.getTypePtr();
+      if (!t) continue;
+      const clang::ElaboratedType *et = t->getAs<clang::ElaboratedType>();
+      if (!et) continue;
+      const clang::NestedNameSpecifier *nns = et->getQualifier();
+      if (!nns) continue;
+      const clang::NamespaceDecl *nsd = nns->getAsNamespace();
+      if (!nsd) continue;
+      const clang::SourceRange sr = valueDecl->getSourceRange();
+      clang::SourceLocation sl(sr.getBegin());
+      clang::SourceLocation end(sr.getEnd());
+      auto &SM = getCompilerInstance().getSourceManager();
+      size_t length = SM.getCharacterData(end) - SM.getCharacterData(sl);
+      StringRef sfull = StringRef(SM.getCharacterData(sl), length);
+      std::string name = nsd->getDeclName().getAsString();
+      size_t offset = sfull.find(name);
+      if (offset > 0) {
+        sl = sl.getLocWithOffset(offset);
+      }
+      FindAndReplace(name, sl, CUDA_CUB_TYPE_NAME_MAP);
+      ret = true;
+    }
+    return ret;
+  }
+  return false;
+}
+
 bool HipifyAction::cudaHostFuncCall(const mat::MatchFinder::MatchResult &Result) {
  if (auto *call = Result.Nodes.getNodeAs<clang::CallExpr>(sCudaHostFuncCall)) {
    if (!call->getNumArgs()) return false;
@@ -555,6 +591,13 @@ std::unique_ptr<clang::ASTConsumer> HipifyAction::CreateASTConsumer(clang::Compi
    ).bind(sCubNamespacePrefix),
    this
  );
+  // TODO: Maybe worth to make it more concrete based on final cubFunctionTemplateDecl
+  Finder->addMatcher(
+    mat::functionTemplateDecl(
+      mat::isExpansionInMainFile()
+    ).bind(sCubFunctionTemplateDecl),
+    this
+  );
  // Ownership is transferred to the caller.
  return Finder->newASTConsumer();
 }
@@ -668,4 +711,5 @@ void HipifyAction::run(const mat::MatchFinder::MatchResult &Result) {
  if (cudaHostFuncCall(Result)) return;
  if (cudaDeviceFuncCall(Result)) return;
  if (cubNamespacePrefix(Result)) return;
+  if (cubFunctionTemplateDecl(Result)) return;
 }
@@ -73,6 +73,7 @@ public:
  bool cudaDeviceFuncCall(const mat::MatchFinder::MatchResult &Result);
  bool cudaHostFuncCall(const mat::MatchFinder::MatchResult &Result);
  bool cubNamespacePrefix(const mat::MatchFinder::MatchResult &Result);
+  bool cubFunctionTemplateDecl(const mat::MatchFinder::MatchResult &Result);
  // Called by the preprocessor for each include directive during the non-raw lexing pass.
  void InclusionDirective(clang::SourceLocation hash_loc,
                          const clang::Token &include_token,
@@ -0,0 +1,70 @@
+// RUN: %run_test hipify "%s" "%t" %hipify_args %clang_args
+// CHECK: #include <hip/hip_runtime.h>
+#include <iostream>
+// CHECK: #include <hiprand.h>
+#include <curand.h>
+// CHECK: #include <hipcub/hipcub.hpp>
+#include <cub/cub.cuh>
+
+#include <iostream>
+
+template <int BLOCK_WIDTH, int ITEMS_PER_THREAD,
+          // CHECK: hipcub::BlockLoadAlgorithm BLOCK_LOAD_ALGO,
+          cub::BlockLoadAlgorithm BLOCK_LOAD_ALGO,
+          // CHECK: hipcub::BlockStoreAlgorithm BLOCK_STORE_ALGO,
+          cub::BlockStoreAlgorithm BLOCK_STORE_ALGO,
+          typename T>
+__global__ void sort(const T* data_in, T* data_out){
+    // CHECK: typedef hipcub::BlockLoad<T, BLOCK_WIDTH, ITEMS_PER_THREAD, BLOCK_LOAD_ALGO> BlockLoadT;
+    typedef cub::BlockLoad<T, BLOCK_WIDTH, ITEMS_PER_THREAD, BLOCK_LOAD_ALGO> BlockLoadT;
+    // CHECK: typedef hipcub::BlockRadixSort<T, BLOCK_WIDTH, ITEMS_PER_THREAD> BlockRadixSortT;
+    typedef cub::BlockRadixSort<T, BLOCK_WIDTH, ITEMS_PER_THREAD> BlockRadixSortT;
+    // CHECK: typedef hipcub::BlockStore<T, BLOCK_WIDTH, ITEMS_PER_THREAD, BLOCK_STORE_ALGO> BlockStoreT;
+    typedef cub::BlockStore<T, BLOCK_WIDTH, ITEMS_PER_THREAD, BLOCK_STORE_ALGO> BlockStoreT;
+    __shared__ union {
+        typename BlockLoadT::TempStorage load;
+        typename BlockRadixSortT::TempStorage sort;
+        typename BlockStoreT::TempStorage store;
+    } tmp_storage;
+    T items[ITEMS_PER_THREAD];
+    BlockLoadT(tmp_storage.load).Load(data_in + blockIdx.x * BLOCK_WIDTH * ITEMS_PER_THREAD, items);
+    __syncthreads();
+    BlockRadixSortT(tmp_storage.sort).Sort(items);
+    __syncthreads();
+    BlockStoreT(tmp_storage.store).Store(data_out + blockIdx.x * BLOCK_WIDTH * ITEMS_PER_THREAD, items);
+}
+
+int main() {
+    double* d_gpu = NULL;
+    double* result_gpu = NULL;
+    double* data_sorted = new double[1000*4096];
+    // Allocate memory on the GPU
+    // CHECK: hipMalloc(&d_gpu, 1000*4096 * sizeof(double));
+    cudaMalloc(&d_gpu, 1000*4096 * sizeof(double));
+    // CHECK: hipMalloc(&result_gpu, 1000*4096 * sizeof(double));
+    cudaMalloc(&result_gpu, 1000*4096 * sizeof(double));
+    // CHECK: hiprandGenerator_t gen;
+    curandGenerator_t gen;
+    // Create generator
+    // CHECK: hiprandCreateGenerator(&gen, HIPRAND_RNG_PSEUDO_DEFAULT);
+    curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT);
+    // Fill array with random numbers
+    // CHECK: hiprandGenerateNormalDouble(gen, d_gpu, 1000*4096, 0.0, 1.0);
+    curandGenerateNormalDouble(gen, d_gpu, 1000*4096, 0.0, 1.0);
+    // Destroy generator
+    // CHECK: hiprandDestroyGenerator(gen);
+    curandDestroyGenerator(gen);
+    // Sort data
+    // TODO: Substitution of cub namespace in CUDAKernelCallExpr
+    // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(sort<512, 8, cub::BLOCK_LOAD_TRANSPOSE, cub::BLOCK_STORE_TRANSPOSE>), dim3(1000), dim3(512), 0, 0, d_gpu, result_gpu);
+    sort<512, 8, cub::BLOCK_LOAD_TRANSPOSE, cub::BLOCK_STORE_TRANSPOSE><<<1000, 512>>>(d_gpu, result_gpu);
+    // CHECK: hipLaunchKernelGGL(HIP_KERNEL_NAME(sort<256, 16, cub::BLOCK_LOAD_DIRECT, cub::BLOCK_STORE_DIRECT>), dim3(1000), dim3(256), 0, 0, d_gpu, result_gpu);
+    sort<256, 16, cub::BLOCK_LOAD_DIRECT, cub::BLOCK_STORE_DIRECT><<<1000, 256>>>(d_gpu, result_gpu);
+    // CHECK: hipMemcpy(data_sorted, result_gpu, 1000*4096*sizeof(double), hipMemcpyDeviceToHost);
+    cudaMemcpy(data_sorted, result_gpu, 1000*4096*sizeof(double), cudaMemcpyDeviceToHost);
+    // Write the sorted data to standard out
+    for (int i = 0; i < 4095; ++i) {
+        std::cout << data_sorted[i] << ", ";
+    }
+    std::cout << data_sorted[4095] << std::endl;
+}