[SYCLomatic] Enable the migration for 4 CUB API BlockShuffle::Offset/Rotate/Up/Down with new introduced helper class group_shuffle(#2613)

intwanghao · web-flow · commit 2d33dba71c73 · 2025-01-20T15:49:01.000+08:00
Signed-off-by: intwanghao &lt;hao3.wang@intel.com&gt;
diff --git a/clang/lib/DPCT/AnalysisInfo.cpp b/clang/lib/DPCT/AnalysisInfo.cpp
@@ -3812,6 +3812,7 @@ void TempStorageVarInfo::addAccessorDecl(StmtList &AccessorList,
     OS << '(' << LocalSize << ".size() * sizeof("
        << ValueType->getSourceString() << ')' << ')';
     break;
+  case BlockShuffle:
   case BlockRadixSort:
     OS << MapNames::getClNamespace() << "local_accessor<uint8_t, 1> " << Name
        << "_acc(";
@@ -3831,6 +3832,7 @@ ParameterStream &TempStorageVarInfo::getFuncDecl(ParameterStream &PS) {
   case BlockReduce:
     PS << MapNames::getClNamespace() << "local_accessor<std::byte, 1> ";
     break;
+  case BlockShuffle:
   case BlockRadixSort:
     PS << "uint8_t *";
     break;
diff --git a/clang/lib/DPCT/AnalysisInfo.h b/clang/lib/DPCT/AnalysisInfo.h
@@ -2237,6 +2237,7 @@ class TempStorageVarInfo {
   enum APIKind {
     BlockReduce,
     BlockRadixSort,
+    BlockShuffle,
   };
 
 private:
diff --git a/clang/lib/DPCT/RuleInfra/APINamesTemplateType.inc b/clang/lib/DPCT/RuleInfra/APINamesTemplateType.inc
@@ -131,6 +131,31 @@ TYPE_REWRITE_ENTRY(
             TYPE_FACTORY(STR(MapNames::getLibraryHelperNamespace() +
                              "group::exchange"),
                          TEMPLATE_ARG(0), TEMPLATE_ARG(2)))))
+
+// cub::BlockShuffle
+TYPE_REWRITE_ENTRY(
+    "cub::BlockShuffle",
+    TYPE_CONDITIONAL_FACTORY(
+        UseSYCLCompat(),
+        WARNING_FACTORY(Diagnostics::UNSUPPORT_SYCLCOMPAT, TYPESTR),
+        HEADER_INSERTION_FACTORY(
+            HeaderType::HT_DPCT_GROUP_Utils,
+            TYPE_CONDITIONAL_FACTORY(
+                CheckTemplateArgCount(2, false),
+                TYPE_FACTORY(STR(MapNames::getLibraryHelperNamespace() +
+                                 "group::group_shuffle"),
+                             TEMPLATE_ARG(0), TEMPLATE_ARG(1)),
+                TYPE_CONDITIONAL_FACTORY(
+                    CheckTemplateArgCount(3, false),
+                    TYPE_FACTORY(STR(MapNames::getLibraryHelperNamespace() +
+                                     "group::group_shuffle"),
+                                 TEMPLATE_ARG(0), TEMPLATE_ARG(1),
+                                 TEMPLATE_ARG(2)),
+                    TYPE_FACTORY(STR(MapNames::getLibraryHelperNamespace() +
+                                     "group::group_shuffle"),
+                                 TEMPLATE_ARG(0), TEMPLATE_ARG(1),
+                                 TEMPLATE_ARG(2), TEMPLATE_ARG(3)))))))
+
 // cub::BlockLoad
 TYPE_REWRITE_ENTRY(
     "cub::BlockLoad",
diff --git a/clang/lib/DPCT/RuleInfra/TypeLocRewriters.cpp b/clang/lib/DPCT/RuleInfra/TypeLocRewriters.cpp
@@ -102,12 +102,24 @@ makeUserDefinedTypeStrCreator(MetaRuleObject &R,
 
 class CheckTemplateArgCount {
   unsigned Count;
+  bool IsIncludeDefault;
 
 public:
-  CheckTemplateArgCount(unsigned I) : Count(I) {}
+  CheckTemplateArgCount(unsigned I, bool D = true)
+      : Count(I), IsIncludeDefault(D) {}
   bool operator()(const TypeLoc TL) {
-    if(auto TSTL = TL.getAs<TemplateSpecializationTypeLoc>()){
-      return TSTL.getNumArgs() == Count;
+    if (auto TSTL = TL.getAs<TemplateSpecializationTypeLoc>()) {
+      size_t Num = TSTL.getNumArgs();
+      if (IsIncludeDefault) {
+        return Num == Count;
+      }
+      size_t NoneDefaultNum = 0;
+      for (int i = 0; i < Num; i++) {
+        if (!TSTL.getArgLoc(i).getArgument().getIsDefaulted()) {
+          NoneDefaultNum++;
+        }
+      }
+      return NoneDefaultNum == Count;
     }
     return false;
   }
diff --git a/clang/lib/DPCT/RulesLang/RewriterSYCLcompat.cpp b/clang/lib/DPCT/RulesLang/RewriterSYCLcompat.cpp
@@ -94,6 +94,10 @@ SYCLCOMPAT_UNSUPPORT("cub::BlockExchange.BlockedToStriped")
 SYCLCOMPAT_UNSUPPORT("cub::BlockExchange.StripedToBlocked")
 SYCLCOMPAT_UNSUPPORT("cub::BlockExchange.ScatterToBlocked")
 SYCLCOMPAT_UNSUPPORT("cub::BlockExchange.ScatterToStriped")
+SYCLCOMPAT_UNSUPPORT("cub::BlockShuffle.Offset")
+SYCLCOMPAT_UNSUPPORT("cub::BlockShuffle.Rotate")
+SYCLCOMPAT_UNSUPPORT("cub::BlockShuffle.Up")
+SYCLCOMPAT_UNSUPPORT("cub::BlockShuffle.Down")
 SYCLCOMPAT_UNSUPPORT("cub::BlockLoad.Load")
 SYCLCOMPAT_UNSUPPORT("cub::BlockStore.Store")
   });
diff --git a/clang/lib/DPCT/RulesLangLib/CUB/RewriterClassMethods.cpp b/clang/lib/DPCT/RulesLangLib/CUB/RewriterClassMethods.cpp
@@ -211,6 +211,66 @@ RewriterMap dpct::createClassMethodsRewriterMap() {
                                 "cub::BlockExchange.ScatterToStriped",
                                 MemberExprBase(), false, "scatter_to_striped",
                                 NDITEM, ARG(0), ARG(1)))
+      // cub::BlockShuffle.Offset
+      HEADER_INSERT_FACTORY(
+          HeaderType::HT_DPCT_GROUP_Utils,
+          CASE_FACTORY_ENTRY(
+              CASE(CheckArgCount(2, std::equal_to<>(), false),
+                   MEMBER_CALL_FACTORY_ENTRY("cub::BlockShuffle.Offset",
+                                             MemberExprBase(), false, "select",
+                                             NDITEM, ARG(0), ARG(1))),
+              CASE(CheckArgCount(3, std::equal_to<>(), false),
+                   MEMBER_CALL_FACTORY_ENTRY("cub::BlockShuffle.Offset",
+                                             MemberExprBase(), false, "select",
+                                             NDITEM, ARG(0), ARG(1), ARG(2))),
+              OTHERWISE(UNSUPPORT_FACTORY_ENTRY("cub::BlockShuffle.Offset",
+                                                Diagnostics::API_NOT_MIGRATED,
+                                                printCallExprPretty()))))
+      // cub::BlockShuffle.Rotate
+      HEADER_INSERT_FACTORY(
+          HeaderType::HT_DPCT_GROUP_Utils,
+          CASE_FACTORY_ENTRY(
+              CASE(CheckArgCount(2, std::equal_to<>(), false),
+                   MEMBER_CALL_FACTORY_ENTRY("cub::BlockShuffle.Rotate",
+                                             MemberExprBase(), false, "select2",
+                                             NDITEM, ARG(0), ARG(1))),
+              CASE(CheckArgCount(3, std::equal_to<>(), false),
+                   MEMBER_CALL_FACTORY_ENTRY("cub::BlockShuffle.Rotate",
+                                             MemberExprBase(), false, "select2",
+                                             NDITEM, ARG(0), ARG(1), ARG(2))),
+              OTHERWISE(UNSUPPORT_FACTORY_ENTRY("cub::BlockShuffle.Rotate",
+                                                Diagnostics::API_NOT_MIGRATED,
+                                                printCallExprPretty()))))
+      // cub::BlockShuffle.Up
+      HEADER_INSERT_FACTORY(
+          HeaderType::HT_DPCT_GROUP_Utils,
+          CASE_FACTORY_ENTRY(
+              CASE(CheckArgCount(2),
+                   MEMBER_CALL_FACTORY_ENTRY(
+                       "cub::BlockShuffle.Up", MemberExprBase(), false,
+                       "shuffle_right", NDITEM, ARG(0), ARG(1))),
+              CASE(CheckArgCount(3),
+                   MEMBER_CALL_FACTORY_ENTRY(
+                       "cub::BlockShuffle.Up", MemberExprBase(), false,
+                       "shuffle_right", NDITEM, ARG(0), ARG(1), ARG(2))),
+              OTHERWISE(UNSUPPORT_FACTORY_ENTRY("cub::BlockShuffle.Up",
+                                                Diagnostics::API_NOT_MIGRATED,
+                                                printCallExprPretty()))))
+      // cub::BlockShuffle.Down
+      HEADER_INSERT_FACTORY(
+          HeaderType::HT_DPCT_GROUP_Utils,
+          CASE_FACTORY_ENTRY(
+              CASE(CheckArgCount(2),
+                   MEMBER_CALL_FACTORY_ENTRY(
+                       "cub::BlockShuffle.Down", MemberExprBase(), false,
+                       "shuffle_left", NDITEM, ARG(0), ARG(1))),
+              CASE(CheckArgCount(3),
+                   MEMBER_CALL_FACTORY_ENTRY(
+                       "cub::BlockShuffle.Down", MemberExprBase(), false,
+                       "shuffle_left", NDITEM, ARG(0), ARG(1), ARG(2))),
+              OTHERWISE(UNSUPPORT_FACTORY_ENTRY("cub::BlockShuffle.Down",
+                                                Diagnostics::API_NOT_MIGRATED,
+                                                printCallExprPretty()))))
       // cub::BlockLoad.Load
       HEADER_INSERT_FACTORY(
           HeaderType::HT_DPCT_GROUP_Utils,
diff --git a/clang/lib/DPCT/RulesLangLib/CUBAPIMigration.cpp b/clang/lib/DPCT/RulesLangLib/CUBAPIMigration.cpp
@@ -100,7 +100,7 @@ void CubTypeRule::registerMatcher(ast_matchers::MatchFinder &MF) {
         "cub::ArgIndexInputIterator", "cub::DiscardOutputIterator",
         "cub::DoubleBuffer", "cub::NullType", "cub::ArgMax", "cub::ArgMin",
         "cub::BlockRadixSort", "cub::BlockExchange", "cub::BlockLoad",
-        "cub::BlockStore");
+        "cub::BlockStore", "cub::BlockShuffle");
   };
 
   MF.addMatcher(
@@ -158,15 +158,16 @@ void CubDeviceLevelRule::runRule(
 void CubMemberCallRule::registerMatcher(ast_matchers::MatchFinder &MF) {
   MF.addMatcher(
       cxxMemberCallExpr(
-          allOf(on(hasType(hasCanonicalType(qualType(hasDeclaration(namedDecl(
-                    hasAnyName("cub::ArgIndexInputIterator",
-                               "cub::BlockRadixSort", "cub::BlockExchange",
-                               "cub::BlockLoad", "cub::BlockStore"))))))),
+          allOf(on(hasType(hasCanonicalType(
+                    qualType(hasDeclaration(namedDecl(hasAnyName(
+                        "cub::ArgIndexInputIterator", "cub::BlockRadixSort",
+                        "cub::BlockExchange", "cub::BlockLoad",
+                        "cub::BlockStore", "cub::BlockShuffle"))))))),
                 callee(cxxMethodDecl(hasAnyName(
                     "normalize", "Sort", "SortDescending", "BlockedToStriped",
                     "StripedToBlocked", "ScatterToBlocked", "ScatterToStriped",
                     "SortBlockedToStriped", "SortDescendingBlockedToStriped",
-                    "Load", "Store")))))
+                    "Load", "Store", "Offset", "Rotate", "Up", "Down")))))
           .bind("memberCall"),
       this);
 
@@ -253,13 +254,17 @@ void CubMemberCallRule::runRule(
         Name == "BlockedToStriped" || Name == "StripedToBlocked" ||
         Name == "StripedToBlocked" || Name == "ScatterToBlocked" ||
         Name == "ScatterToStriped";
-    if (isBlockRadixSort || isBlockExchange || Name == "Load" ||
-        Name == "Store") {
+    bool isBlockShuffle =
+        Name == "Offset" || Name == "Rotate" || Name == "Up" || Name == "Down";
+    if (isBlockRadixSort || isBlockExchange || isBlockShuffle ||
+        Name == "Load" || Name == "Store") {
       std::string HelpFuncName;
       if (isBlockRadixSort)
         HelpFuncName = "group_radix_sort";
       else if (isBlockExchange)
         HelpFuncName = "exchange";
+      else if (isBlockShuffle)
+        HelpFuncName = "group_shuffle";
       else if (Name == "Load")
         HelpFuncName = "group_load";
       else if (Name == "Store")
@@ -273,20 +278,36 @@ void CubMemberCallRule::runRule(
       auto *ClassSpecDecl = dyn_cast<ClassTemplateSpecializationDecl>(
           CanTy->getAs<RecordType>()->getDecl());
       const auto &ValueTyArg = ClassSpecDecl->getTemplateArgs()[0];
-      const auto &ItemsPreThreadArg = ClassSpecDecl->getTemplateArgs()[2];
+
       ValueTyArg.getAsType().getAsString();
       std::string Fn;
       llvm::raw_string_ostream OS(Fn);
       OS << MapNames::getDpctNamespace() << "group::" << HelpFuncName << "<"
-         << ValueTyArg.getAsType().getAsString() << ", "
-         << ItemsPreThreadArg.getAsIntegral() << ">::get_local_memory_size";
+         << ValueTyArg.getAsType().getAsString();
+      if (isBlockShuffle) {
+        if (!ClassSpecDecl->getTemplateArgs()[1].getIsDefaulted()) {
+          OS << ", " << ClassSpecDecl->getTemplateArgs()[1].getAsIntegral();
+        }
+        if (!ClassSpecDecl->getTemplateArgs()[2].getIsDefaulted()) {
+          OS << ", " << ClassSpecDecl->getTemplateArgs()[2].getAsIntegral();
+        }
+        if (!ClassSpecDecl->getTemplateArgs()[3].getIsDefaulted()) {
+          OS << ", " << ClassSpecDecl->getTemplateArgs()[3].getAsIntegral();
+        }
+      } else {
+        const auto &ItemsPreThreadArg = ClassSpecDecl->getTemplateArgs()[2];
+        OS << ", " << ItemsPreThreadArg.getAsIntegral();
+      }
+      OS << ">::get_local_memory_size";
       if (auto FuncInfo = DeviceFunctionDecl::LinkRedecls(FD)) {
         auto LocInfo = DpctGlobalInfo::getLocInfo(TempStorage);
         ExprAnalysis EA;
         EA.analyze(DataTypeLoc);
         FuncInfo->getVarMap().addCUBTempStorage(
             std::make_shared<TempStorageVarInfo>(
-                LocInfo.second, TempStorageVarInfo::BlockRadixSort,
+                LocInfo.second,
+                isBlockShuffle ? TempStorageVarInfo::BlockShuffle
+                               : TempStorageVarInfo::BlockRadixSort,
                 TempStorage->getName(), Fn,
                 EA.getTemplateDependentStringInfo()));
       }
diff --git a/clang/lib/DPCT/SrcAPI/APINames_CUB.inc b/clang/lib/DPCT/SrcAPI/APINames_CUB.inc
@@ -119,10 +119,10 @@ ENTRY_MEMBER_FUNCTION(cub::BlockScan, cub::BlockScan, ExclusiveSum, ExclusiveSum
 ENTRY_MEMBER_FUNCTION(cub::BlockScan, cub::BlockScan, ExclusiveScan, ExclusiveScan, true, NO_FLAG, P4, "Successful")
 ENTRY_MEMBER_FUNCTION(cub::BlockScan, cub::BlockScan, InclusiveSum, InclusiveSum, true, NO_FLAG, P4, "Successful")
 ENTRY_MEMBER_FUNCTION(cub::BlockScan, cub::BlockScan, InclusiveScan, InclusiveScan, true, NO_FLAG, P4, "Successful")
-ENTRY_MEMBER_FUNCTION(cub::BlockShuffle, cub::BlockShuffle, Offset, Offset, false, NO_FLAG, P4, "Comment")
-ENTRY_MEMBER_FUNCTION(cub::BlockShuffle, cub::BlockShuffle, Rotate, Rotate, false, NO_FLAG, P4, "Comment")
-ENTRY_MEMBER_FUNCTION(cub::BlockShuffle, cub::BlockShuffle, Up, Up, false, NO_FLAG, P4, "Comment")
-ENTRY_MEMBER_FUNCTION(cub::BlockShuffle, cub::BlockShuffle, Down, Down, false, NO_FLAG, P4, "Comment")
+ENTRY_MEMBER_FUNCTION(cub::BlockShuffle, cub::BlockShuffle, Offset, Offset, true, NO_FLAG, P4, "Comment")
+ENTRY_MEMBER_FUNCTION(cub::BlockShuffle, cub::BlockShuffle, Rotate, Rotate, true, NO_FLAG, P4, "Comment")
+ENTRY_MEMBER_FUNCTION(cub::BlockShuffle, cub::BlockShuffle, Up, Up, true, NO_FLAG, P4, "Comment")
+ENTRY_MEMBER_FUNCTION(cub::BlockShuffle, cub::BlockShuffle, Down, Down, true, NO_FLAG, P4, "Comment")
 ENTRY_MEMBER_FUNCTION(cub::BlockDiscontinuity, cub::BlockDiscontinuity, FlagHeads, FlagHeads, false, NO_FLAG, P4, "Comment")
 ENTRY_MEMBER_FUNCTION(cub::BlockDiscontinuity, cub::BlockDiscontinuity, FlagTails, FlagTails, false, NO_FLAG, P4, "Comment")
 ENTRY_MEMBER_FUNCTION(cub::BlockDiscontinuity, cub::BlockDiscontinuity, FlagHeadsAndTails, FlagHeadsAndTails, false, NO_FLAG, P4, "Comment")
diff --git a/clang/runtime/dpct-rt/include/dpct/group_utils.hpp b/clang/runtime/dpct-rt/include/dpct/group_utils.hpp
diff --git a/clang/test/dpct/cub/blocklevel/blockshuffle.cu b/clang/test/dpct/cub/blocklevel/blockshuffle.cu
