#1: fix more bugs in tag output

DARMA-tasking · Feb 16, 2021 · 15bd33a · 15bd33a
1 parent 5f20883
commit 15bd33a
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 14 deletions.
diff --git a/src/test.cc b/src/test.cc
@@ -13,10 +13,13 @@ template <typename... A>
 class RangePolicy {
 public:
   explicit RangePolicy(size_t in) {}
+  explicit RangePolicy(size_t in, size_t out) {}
 };
 
 template <typename Policy>
 void parallel_for(std::string str, Policy rp) { }
+template <typename Policy, typename X>
+void parallel_for(std::string str, Policy rp, X) { }
 
 void parallel_for(std::string str, size_t x) { }
 void parallel_for(size_t x) { }
@@ -26,25 +29,43 @@ void fence() { }
 
 }
 
+namespace PHX {
+class Device {};
+}
+
 void test3() {
-  Kokkos::parallel_for("test-123", Kokkos::RangePolicy<int, float>(1000));
+  empire::parallel_for_async("test-123", buildKokkosPolicy<Kokkos::RangePolicy, int, float>(1000));
 }
 
 void test() {
   //using namespace Kokkos;
   while (true) {
     int x = 10;
     if (true) x=5;
-    Kokkos::parallel_for("80", 90);
-    Kokkos::fence();
+    empire::parallel_for_blocking("80", 90);
     if (false) x=6;
   }
 }
 
+#define KOKKOS_LAMBDA []
+
+struct X {
+  size_t extent(int len){return 0;};
+};
+
 template <typename T>
 void test4() {
-  Kokkos::parallel_for("10", 20);
-  Kokkos::fence();
+  using Size = size_t;
+
+  empire::parallel_for_blocking("10", 20);
+
+  X S_values;
+
+  empire::parallel_for_blocking("hello", buildKokkosPolicy<Kokkos::RangePolicy, PHX::Device>(0,S_values.extent(0)),
+  KOKKOS_LAMBDA (const Size& n)
+  {
+    return n+1;
+  });
 }
 
 void test2() {
@@ -53,17 +74,14 @@ void test2() {
   test4<long>();
 
   if (true) {
-    Kokkos::parallel_for("10", 20);
-    Kokkos::fence();
-    Kokkos::parallel_scan("10", 20);
-    Kokkos::fence();
+    empire::parallel_for_blocking("10", 20);
+    empire::parallel_scan_blocking("10", 20);
   }
   if (true) {
-    Kokkos::parallel_for(40);
-    Kokkos::fence();
+    empire::parallel_for_blocking(40);
   }
   if (true) {
-    Kokkos::parallel_for("1219", 23);
+    empire::parallel_for_async("1219", 23);
     test();
   }
 

diff --git a/src/transform.cc b/src/transform.cc
@@ -165,15 +165,21 @@ struct RewriteArgument {
       // advanced policy must lift to builder
       fmt::print("range is advanced policy\n");
       //policy_type.dump();
-      auto str = policy_type.getAsString();
+      int offset = 0;
+      PrintingPolicy p(LangOptions{});
+      p.SuppressTagKeyword = true;
+      auto str = policy_type.getAsString(p);
       if (str.substr(0, 6) == "const ") {
         str = str.substr(6, str.size()-1);
+        offset += 6;
       }
       if (str.substr(0, 6) == "class ") {
         str = str.substr(6, str.size()-1);
+        offset += 6;
       }
       if (str.substr(0, 7) == "struct ") {
         str = str.substr(7, str.size()-1);
+        offset += 7;
       }
 
       auto t = str;
@@ -203,7 +209,7 @@ struct RewriteArgument {
         new_type = fmt::format("buildKokkosPolicy<{}, {}>", main_type, temp_type);
       }
       fmt::print("new type={}\n", new_type);
-      rw.ReplaceText(begin, str.size()-6, new_type);
+      rw.ReplaceText(begin, t.size(), new_type);
     }
   }