[SYCLomatic] Support migration of cub::{LoadDirectBlocked, LoadDirectStriped} (#2094)

Signed-off-by: Wang, Yihan <[email protected]>

Author: Yihan Wang

clang/lib/DPCT/APINames_CUB.inc

@@ -228,11 +228,11 @@ ENTRY_MEMBER_FUNCTION(cub::GridEvenShare, cub::GridEvenShare, BlockInit, BlockIn
 // Thread and thread block I/O
 ENTRY(cub::ThreadLoad, cub::ThreadLoad, true, NO_FLAG, P4, "Successful")
 ENTRY(cub::ThreadStore, cub::ThreadStore, true, NO_FLAG, P4, "Successful")
-ENTRY(cub::LoadDirectBlocked, cub::LoadDirectBlocked, false, NO_FLAG, P4, "Comment")
+ENTRY(cub::LoadDirectBlocked, cub::LoadDirectBlocked, true, NO_FLAG, P4, "Successful")
 ENTRY(cub::LoadDirectBlockedVectorized, cub::LoadDirectBlockedVectorized, false, NO_FLAG, P4, "Comment")
 ENTRY(cub::StoreDirectBlocked, cub::StoreDirectBlocked, false, NO_FLAG, P4, "Comment")
 ENTRY(cub::StoreDirectBlockedVectorized, cub::StoreDirectBlockedVectorized, false, NO_FLAG, P4, "Comment")
-ENTRY(cub::LoadDirectStriped, cub::LoadDirectStriped, false, NO_FLAG, P4, "Comment")
+ENTRY(cub::LoadDirectStriped, cub::LoadDirectStriped, true, NO_FLAG, P4, "Successful")
 ENTRY(cub::StoreDirectStriped, cub::StoreDirectStriped, false, NO_FLAG, P4, "Comment")
 ENTRY(cub::LoadDirectWarpStriped, cub::LoadDirectWarpStriped, false, NO_FLAG, P4, "Comment")
 ENTRY(cub::StoreDirectWarpStriped, cub::StoreDirectWarpStriped, false, NO_FLAG, P4, "Comment")

clang/lib/DPCT/CUBAPIMigration.cpp

@@ -192,7 +192,8 @@ void CubIntrinsicRule::registerMatcher(ast_matchers::MatchFinder &MF) {
                          "WarpId", "SyncStream", "CurrentDevice", "DeviceCount",
                          "DeviceCountUncached", "DeviceCountCachedValue",
                          "PtxVersion", "PtxVersionUncached", "SmVersion",
-                         "SmVersionUncached", "RowMajorTid"),
+                         "SmVersionUncached", "RowMajorTid",
+                         "LoadDirectBlocked", "LoadDirectStriped"),
               hasAncestor(namespaceDecl(hasName("cub")))))))
           .bind("IntrinsicCall"),
       this);

clang/lib/DPCT/Rewriters/CUB/RewriterUtilityFunctions.cpp

@@ -114,5 +114,19 @@ RewriterMap dpct::createUtilityFunctionsRewriterMap() {
                 LITERAL("10"))))
       // cub::RowMajorTid
       MEMBER_CALL_FACTORY_ENTRY("cub::RowMajorTid", NDITEM, /*IsArrow=*/false,
-                                "get_local_linear_id")};
+                                "get_local_linear_id")
+      // cub::LoadDirectBlocked
+      HEADER_INSERT_FACTORY(
+          HeaderType::HT_DPCT_GROUP_Utils,
+          CALL_FACTORY_ENTRY(
+              "cub::LoadDirectBlocked",
+              CALL(MapNames::getDpctNamespace() + "group::load_direct_blocked",
+                   NDITEM, ARG(1), ARG(2))))
+      // cub::LoadDirectStriped
+      HEADER_INSERT_FACTORY(
+          HeaderType::HT_DPCT_GROUP_Utils,
+          CALL_FACTORY_ENTRY(
+              "cub::LoadDirectStriped",
+              CALL(MapNames::getDpctNamespace() + "group::load_direct_striped",
+                   NDITEM, ARG(1), ARG(2))))};
 }

clang/runtime/dpct-rt/include/dpct/group_utils.hpp

@@ -476,6 +476,41 @@ __dpct_inline__ void load_striped(const Item &item, InputIteratorT block_itr,
   }
 }
 
+// loads a linear segment of workgroup items into a blocked arrangement.
+template <typename InputT, size_t ITEMS_PER_WORK_ITEM, typename InputIteratorT,
+          typename Item>
+__dpct_inline__ void load_direct_blocked(const Item &item, InputIteratorT block_itr,
+                                  InputT (&items)[ITEMS_PER_WORK_ITEM]) {
+
+  // This implementation does not take in account range loading across
+  // workgroup items To-do: Decide whether range loading is required for group
+  // loading
+  size_t linear_tid = item.get_local_linear_id();
+  uint32_t workgroup_offset = linear_tid * ITEMS_PER_WORK_ITEM;
+#pragma unroll
+  for (size_t idx = 0; idx < ITEMS_PER_WORK_ITEM; idx++) {
+    items[idx] = block_itr[workgroup_offset + idx];
+  }
+}
+
+// loads a linear segment of workgroup items into a striped arrangement.
+template <typename InputT, size_t ITEMS_PER_WORK_ITEM, typename InputIteratorT,
+          typename Item>
+__dpct_inline__ void load_direct_striped(const Item &item, InputIteratorT block_itr,
+                                  InputT (&items)[ITEMS_PER_WORK_ITEM]) {
+
+  // This implementation does not take in account range loading across
+  // workgroup items To-do: Decide whether range loading is required for group
+  // loading
+  size_t linear_tid = item.get_local_linear_id();
+  size_t group_work_items = item.get_local_range().size();
+#pragma unroll
+  for (size_t idx = 0; idx < ITEMS_PER_WORK_ITEM; idx++) {
+    items[idx] = block_itr[linear_tid + (idx * group_work_items)];
+  }
+}
+
+
 // loads a linear segment of workgroup items into a subgroup striped
 // arrangement. Created as free function until exchange mechanism is
 // implemented.

clang/test/dpct/cub/intrinsic/load.cu

@@ -0,0 +1,22 @@
+// UNSUPPORTED: cuda-8.0, cuda-9.0, cuda-9.1, cuda-9.2, cuda-10.0, cuda-10.1, cuda-10.2
+// UNSUPPORTED: v8.0, v9.0, v9.1, v9.2, v10.0, v10.1, v10.2
+// RUN: dpct -in-root %S -out-root %T/intrinsic/load %S/load.cu --cuda-include-path="%cuda-path/include" -- -std=c++14 -x cuda --cuda-host-only
+// RUN: FileCheck --input-file %T/intrinsic/load/load.dp.cpp --match-full-lines %s
+// RUN: %if build_lit %{icpx -c -fsycl %T/intrinsic/load/load.dp.cpp -o %T/intrinsic/load/load.dp.o %}
+
+// CHECK:#include <sycl/sycl.hpp>
+// CHECK:#include <dpct/dpct.hpp>
+// CHECK:#include <dpct/group_utils.hpp>
+#include <cub/cub.cuh>
+
+__global__ void TestLoadStriped(int *d_data) {
+  int thread_data[4];
+  // CHECK: dpct::group::load_direct_striped(item_ct1, d_data, thread_data);
+  cub::LoadDirectStriped<128>(threadIdx.x, d_data, thread_data);
+}
+
+__global__ void BlockedToStripedKernel(int *d_data) {
+  int thread_data[4];
+  // CHECK: dpct::group::load_direct_blocked(item_ct1, d_data, thread_data);
+  cub::LoadDirectBlocked(threadIdx.x, d_data, thread_data);
+}