[SYCLomatic] Support migration of thrust random types (#1002)

Signed-off-by: chenwei.sun <[email protected]>

Author: tomflinda

clang/lib/DPCT/APINamesTemplateType.inc

@@ -200,4 +200,16 @@ TYPE_REWRITE_ENTRY("cooperative_groups::__v1::bit_or",
                    TYPE_FACTORY(STR(MapNames::getClNamespace() + "bit_or"),
                                 TEMPLATE_ARG(0)))
 
+TYPE_REWRITE_ENTRY("thrust::random::default_random_engine",
+                   TYPE_FACTORY(STR("oneapi::dpl::default_engine")))
+
+TYPE_REWRITE_ENTRY("thrust::random::uniform_real_distribution",
+                   TYPE_FACTORY(STR("oneapi::dpl::uniform_real_distribution"), TEMPLATE_ARG(0)))
+
+TYPE_REWRITE_ENTRY("thrust::random::normal_distribution",
+                   TYPE_FACTORY(STR("oneapi::dpl::normal_distribution"), TEMPLATE_ARG(0)))
+
+TYPE_REWRITE_ENTRY("thrust::random::uniform_int_distribution",
+                   TYPE_FACTORY(STR("oneapi::dpl::uniform_int_distribution"), TEMPLATE_ARG(0)))
+
 // clang-format on

clang/lib/DPCT/ASTTraversal.cpp

@@ -1123,6 +1123,8 @@ void IncludesCallbacks::InclusionDirective(
     } else if (FileName.compare(StringRef("thrust/uninitialized_copy.h")) ==
                0) {
       DpctGlobalInfo::getInstance().insertHeader(HashLoc, HT_DPL_Memory);
+    } else if (FileName.compare(StringRef("thrust/random.h")) == 0) {
+      DpctGlobalInfo::getInstance().insertHeader(HashLoc, HT_DPL_Random);
     } else {
       if(FileName.compare(StringRef("thrust/functional.h")) == 0)
         DpctGlobalInfo::getInstance().insertHeader(HashLoc, HT_Functional);
@@ -14644,17 +14646,37 @@ void TemplateSpecializationTypeLocRule::registerMatcher(
     ast_matchers::MatchFinder &MF) {
   auto TargetTypeName = [&]() {
     return hasAnyName("thrust::not_equal_to", "thrust::constant_iterator",
-                      "thrust::system::cuda::experimental::pinned_allocator");
+                      "thrust::system::cuda::experimental::pinned_allocator",
+                      "thrust::random::default_random_engine",
+                      "thrust::random::uniform_real_distribution",
+                      "thrust::random::normal_distribution",
+                      "thrust::random::linear_congruential_engine",
+                      "thrust::random::uniform_int_distribution");
   };
 
-  MF.addMatcher(typeLoc(
-                    loc(qualType(hasDeclaration(namedDecl(TargetTypeName())))))
-                    .bind("loc"),
-                this);
+  MF.addMatcher(
+      typeLoc(loc(qualType(hasDeclaration(namedDecl(TargetTypeName())))))
+          .bind("loc"),
+      this);
+
+  MF.addMatcher(declRefExpr().bind("declRefExpr"), this);
 }
 
 void TemplateSpecializationTypeLocRule::runRule(
     const ast_matchers::MatchFinder::MatchResult &Result) {
+
+  const DeclRefExpr *DRE = getNodeAsType<DeclRefExpr>(Result, "declRefExpr");
+  if (DRE) {
+    std::string TypeName = DpctGlobalInfo::getTypeName(DRE->getType());
+    std::string Name = DRE->getNameInfo().getName().getAsString();
+    if (TypeName.find("thrust::random::linear_congruential_engine") !=
+            std::string::npos &&
+        Name == "max") {
+      emplaceTransformation(
+          new ReplaceStmt(DRE, "oneapi::dpl::default_engine::max()"));
+    }
+  }
+
   if (auto TL = getNodeAsType<TypeLoc>(Result, "loc")) {
     ExprAnalysis EA;
     EA.analyze(*TL);

clang/lib/DPCT/HeaderTypes.inc

@@ -52,6 +52,7 @@ ONEDPL_HEADER(Execution, "<oneapi/dpl/execution>")
 ONEDPL_HEADER(Iterator, "<oneapi/dpl/iterator>")
 ONEDPL_HEADER(Async, "<oneapi/dpl/async>")
 ONEDPL_HEADER(Memory, "<oneapi/dpl/memory>")
+ONEDPL_HEADER(Random, "<oneapi/dpl/random>")
 
 MKL_HEADER(Mkl, "<oneapi/mkl.hpp>")
 MKL_HEADER(RNG, "<oneapi/mkl/rng/device.hpp>")

clang/test/dpct/thrust-random-type.cu

@@ -0,0 +1,67 @@
+// UNSUPPORTED: cuda-8.0
+// UNSUPPORTED: v8.0
+// RUN: dpct --format-range=none -out-root %T/thrust-random-type %s --cuda-include-path="%cuda-path/include" -- -std=c++14 -x cuda --cuda-host-only
+// RUN: FileCheck --input-file %T/thrust-random-type/thrust-random-type.dp.cpp --match-full-lines %s
+
+
+#include <iostream>
+#include <thrust/device_vector.h>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/random.h>
+#include <thrust/random/linear_congruential_engine.h>
+#include <thrust/random/uniform_real_distribution.h>
+#include <thrust/transform.h>
+
+// CHECK:struct random_1 {
+// CHECK-NEXT:  float operator()(const unsigned int n) const {
+// CHECK-NEXT:    oneapi::dpl::default_engine rng;
+// CHECK-NEXT:    oneapi::dpl::uniform_real_distribution<float> dist(1.0f, 2.0f);
+// CHECK-NEXT:    rng.discard(n);
+// CHECK-NEXT:    return dist(rng);
+// CHECK-NEXT:  }
+// CHECK-NEXT:};
+struct random_1 {
+  __host__ __device__ float operator()(const unsigned int n) const {
+    thrust::default_random_engine rng;
+    thrust::uniform_real_distribution<float> dist(1.0f, 2.0f);
+    rng.discard(n);
+    return dist(rng);
+  }
+};
+
+
+// CHECK:struct random_2 {
+// CHECK-NEXT:  float operator()(const unsigned int n) const {
+// CHECK-NEXT:    oneapi::dpl::default_engine rng;
+// CHECK-NEXT:    rng.discard(n);
+// CHECK-NEXT:    return (float)rng() / oneapi::dpl::default_engine::max();
+// CHECK-NEXT:  }
+// CHECK-NEXT:};
+struct random_2 {
+  __device__ float operator()(const unsigned int n) {
+    thrust::default_random_engine rng;
+    rng.discard(n);
+    return (float)rng() / thrust::default_random_engine::max;
+  }
+};
+
+void test(void) {
+  {
+    const int N = 20;
+    // CHECK:    dpct::device_vector<float> numbers(N);
+    // CHECK-NEXT:    oneapi::dpl::counting_iterator<unsigned int> index_sequence_begin(0);
+    // CHECK-NEXT:    std::transform(oneapi::dpl::execution::seq, index_sequence_begin, index_sequence_begin + N, numbers.begin(), random_1());
+    // CHECK-NEXT:    std::transform(oneapi::dpl::execution::seq, index_sequence_begin, index_sequence_begin + N, numbers.begin(), random_2());
+    thrust::device_vector<float> numbers(N);
+    thrust::counting_iterator<unsigned int> index_sequence_begin(0);
+    thrust::transform(index_sequence_begin, index_sequence_begin + N, numbers.begin(), random_1());
+    thrust::transform(index_sequence_begin, index_sequence_begin + N, numbers.begin(), random_2());
+  }
+
+  {
+    // CHECK:    oneapi::dpl::uniform_int_distribution<int> dist1(-5, 10);
+    // CHECK-NEXT:    oneapi::dpl::normal_distribution<float> dist2(1.0f, 2.0f);
+    thrust::uniform_int_distribution<int> dist1(-5, 10);
+    thrust::normal_distribution<float> dist2(1.0f, 2.0f);
+  }
+}