[SYCLomatic] Refine the implementation of dpct::group::load_direct_{blocked, striped} (#2248)

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

Author: Yihan Wang

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

@@ -6,11 +6,53 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "AnalysisInfo.h"
 #include "CallExprRewriterCUB.h"
 #include "CallExprRewriterCommon.h"
 
 using namespace clang::dpct;
 
+namespace {
+class PrettyTemplatedFunctionNamePrinter {
+  std::string Name;
+  std::vector<TemplateArgumentInfo> Args;
+
+public:
+  PrettyTemplatedFunctionNamePrinter(StringRef Name,
+                                     std::vector<TemplateArgumentInfo> &&Args)
+      : Name(Name.str()), Args(std::move(Args)) {}
+  template <class StreamT> void print(StreamT &Stream) const {
+    dpct::print(Stream, Name);
+    if (!Args.empty()) {
+      Stream << '<';
+      ArgsPrinter<false, std::vector<TemplateArgumentInfo>>(Args).print(Stream);
+      Stream << '>';
+    }
+  }
+};
+
+std::function<PrettyTemplatedFunctionNamePrinter(const CallExpr *)>
+makePrettyTemplatedCalleeCreator(std::string CalleeName,
+                                 std::vector<size_t> Indexes) {
+  return PrinterCreator<
+      PrettyTemplatedFunctionNamePrinter, std::string,
+      std::function<std::vector<TemplateArgumentInfo>(const CallExpr *)>>(
+      CalleeName, [=](const CallExpr *C) -> std::vector<TemplateArgumentInfo> {
+        std::vector<TemplateArgumentInfo> Ret;
+        auto List = getTemplateArgsList(C);
+        for (auto Idx : Indexes) {
+          if (Idx < List.size()) {
+            Ret.emplace_back(List[Idx]);
+          }
+        }
+        return Ret;
+      });
+}
+} // namespace
+
+#define PRETTY_TEMPLATED_CALLEE(FuncName, ...)                                 \
+  makePrettyTemplatedCalleeCreator(FuncName, {__VA_ARGS__})
+
 RewriterMap dpct::createUtilityFunctionsRewriterMap() {
   return RewriterMap{
       // cub::IADD3
@@ -120,13 +162,17 @@ RewriterMap dpct::createUtilityFunctionsRewriterMap() {
           HeaderType::HT_DPCT_GROUP_Utils,
           CALL_FACTORY_ENTRY(
               "cub::LoadDirectBlocked",
-              CALL(MapNames::getDpctNamespace() + "group::load_direct_blocked",
-                   NDITEM, ARG(1), ARG(2))))
+              CALL(PRETTY_TEMPLATED_CALLEE(MapNames::getDpctNamespace() +
+                                               "group::load_direct_blocked",
+                                           0, 1, 2),
+                   ARG(0), 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))))};
+              CALL(PRETTY_TEMPLATED_CALLEE(MapNames::getDpctNamespace() +
+                                               "group::load_direct_striped",
+                                           0, 1, 2, 3),
+                   ARG(0), ARG(1), ARG(2))))};
 }

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

@@ -9,6 +9,7 @@
 #ifndef __DPCT_GROUP_UTILS_HPP__
 #define __DPCT_GROUP_UTILS_HPP__
 
+#include <iterator>
 #include <stdexcept>
 #include <sycl/sycl.hpp>
 
@@ -476,41 +477,59 @@ __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;
+/// Load a linear segment of elements into a blocked arrangement across the
+/// work-group.
+///
+/// \tparam InputT The data type to load.
+///
+/// \tparam ElementsPerWorkItem The number of consecutive elements partitioned
+/// onto each work-item.
+///
+/// \tparam InputIteratorT  The random-access iterator type for input \iterator.
+///
+/// \param linear_tid A suitable linear identifier for the calling work-item.
+///
+/// \param block_itr The work-group's base input iterator for loading from.
+///
+/// \param items Data to load
+template <typename InputT, size_t ElementsPerWorkItem, typename InputIteratorT>
+__dpct_inline__ void load_direct_blocked(size_t linear_tid,
+                                         InputIteratorT block_itr,
+                                         InputT (&items)[ElementsPerWorkItem]) {
 #pragma unroll
-  for (size_t idx = 0; idx < ITEMS_PER_WORK_ITEM; idx++) {
-    items[idx] = block_itr[workgroup_offset + idx];
+  for (size_t i = 0; i < ElementsPerWorkItem; i++) {
+    items[i] = block_itr[(linear_tid * ElementsPerWorkItem) + i];
   }
 }
 
-// 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();
+/// Load a linear segment of elements into a striped arrangement across the
+/// work-group.
+///
+/// \tparam WorkGroupSize The work-group size.
+///
+/// \tparam InputT The data type to load.
+///
+/// \tparam ElementsPerWorkItem The number of consecutive elements partitioned
+/// onto each work-item.
+///
+/// \tparam InputIteratorT  The random-access iterator type for input \iterator.
+///
+/// \param linear_tid A suitable linear identifier for the calling work-item.
+///
+/// \param block_itr The work-group's base input iterator for loading from.
+///
+/// \param items Data to load
+template <size_t WorkGroupSize, typename InputT, int ElementsPerWorkItem,
+          typename InputIteratorT>
+__dpct_inline__ void load_direct_striped(size_t linear_tid,
+                                         InputIteratorT block_itr,
+                                         InputT (&items)[ElementsPerWorkItem]) {
 #pragma unroll
-  for (size_t idx = 0; idx < ITEMS_PER_WORK_ITEM; idx++) {
-    items[idx] = block_itr[linear_tid + (idx * group_work_items)];
+  for (size_t i = 0; i < ElementsPerWorkItem; i++) {
+    items[i] = block_itr[linear_tid + i * WorkGroupSize];
   }
 }
 
-
 // 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

@@ -11,12 +11,12 @@
 
 __global__ void TestLoadStriped(int *d_data) {
   int thread_data[4];
-  // CHECK: dpct::group::load_direct_striped(item_ct1, d_data, thread_data);
+  // CHECK: dpct::group::load_direct_striped<128>(item_ct1.get_local_id(2), 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);
+  // CHECK: dpct::group::load_direct_blocked(item_ct1.get_local_id(2), d_data, thread_data);
   cub::LoadDirectBlocked(threadIdx.x, d_data, thread_data);
 }