[SYCLomatic] Refine the migration of cub::ShuffleIndex/Up/Down (#2738)

Signed-off-by: intwanghao <[email protected]>

Author: intwanghao

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

@@ -336,9 +336,11 @@ RewriterMap dpct::createUtilityFunctionsRewriterMap() {
                                            "experimental::shift_sub_group_left",
                                        0, 1),
                       SUBGROUP, ARG(0), ARG(1), ARG(2), ARG(3))),
-              UNSUPPORT_FACTORY_ENTRY("cub::ShuffleDown",
-                                      Diagnostics::API_NOT_MIGRATED,
-                                      LITERAL("cub::ShuffleDown"))))
+              UNSUPPORT_FACTORY_ENTRY(
+                  "cub::ShuffleDown",
+                  Diagnostics::TRY_EXPERIMENTAL_FEATURE,
+                  LITERAL("cub::ShuffleDown"),
+                  LITERAL("--use-experimental-features=non-uniform-groups"))))
       // cub::ShuffleUp
       SUBGROUPSIZE_FACTORY(
           UINT_MAX,
@@ -352,7 +354,8 @@ RewriterMap dpct::createUtilityFunctionsRewriterMap() {
                                "experimental::shift_sub_group_right",
                            0, 1),
                        SUBGROUP, ARG(0), ARG(1), ARG(2), ARG(3))),
-              UNSUPPORT_FACTORY_ENTRY("cub::ShuffleUp",
-                                      Diagnostics::API_NOT_MIGRATED,
-                                      LITERAL("cub::ShuffleUp"))))};
+              UNSUPPORT_FACTORY_ENTRY(
+                  "cub::ShuffleUp", Diagnostics::TRY_EXPERIMENTAL_FEATURE,
+                  LITERAL("cub::ShuffleUp"),
+                  LITERAL("--use-experimental-features=non-uniform-groups"))))};
 }

clang/lib/DPCT/RulesLangLib/CUBAPIMigration.cpp

@@ -1081,38 +1081,43 @@ void CubRule::processWarpLevelFuncCall(const CallExpr *CE, bool FuncCallUsed) {
     if (!TA)
       return;
     WarpSize = TA->get(0).getAsIntegral().getExtValue();
-    std::string ValueType =
-        TA->get(1).getAsType().getUnqualifiedType().getAsString();
     const auto *MemberMask = CE->getArg(2);
-    const auto *Mask = dyn_cast<IntegerLiteral>(MemberMask);
     const Expr *Value = CE->getArg(0);
     const Expr *Lane = CE->getArg(1);
     const auto *DeviceFuncDecl = getImmediateOuterFuncDecl(CE);
     ExprAnalysis ValueEA(Value);
     ExprAnalysis LaneEA(Lane);
     llvm::raw_string_ostream OS(Repl);
-    if (Mask) {
+
+    if (const auto *Mask =
+            dyn_cast<IntegerLiteral>(MemberMask->IgnoreImplicitAsWritten())) {
       if (Mask->getValue().getZExtValue() == 0xffffffff) {
         OS << MapNames::getDpctNamespace() << "select_from_sub_group("
            << DpctGlobalInfo::getSubGroup(CE, DeviceFuncDecl) << ", "
            << ValueEA.getReplacedString() << ", " << LaneEA.getReplacedString();
         if (WarpSize != 32)
           OS << ", " << WarpSize;
         OS << ')';
-      } else {
-        OS << MapNames::getDpctNamespace() << "experimental::"
-           << "select_from_sub_group(" << getStmtSpelling(Mask) << ", "
-           << DpctGlobalInfo::getSubGroup(CE, DeviceFuncDecl) << ", "
-           << ValueEA.getReplacedString() << ", " << LaneEA.getReplacedString();
-        if (WarpSize != 32)
-          OS << ", " << WarpSize;
-        OS << ')';
+        emplaceTransformation(new ReplaceStmt(CE, Repl));
+        return;
       }
+    }
+    if (DpctGlobalInfo::useExpNonUniformGroups()) {
+      ExprAnalysis MaskEA(MemberMask);
+      OS << MapNames::getDpctNamespace() << "experimental::"
+         << "select_from_sub_group(" << MaskEA.getReplacedString() << ", "
+         << DpctGlobalInfo::getSubGroup(CE, DeviceFuncDecl) << ", "
+         << ValueEA.getReplacedString() << ", " << LaneEA.getReplacedString();
+      if (WarpSize != 32)
+        OS << ", " << WarpSize;
+      OS << ')';
       emplaceTransformation(new ReplaceStmt(CE, Repl));
     } else {
-      report(CE->getBeginLoc(), Diagnostics::API_NOT_MIGRATED, false,
-             GetFunctionName(CE));
+      report(CE->getBeginLoc(), Diagnostics::TRY_EXPERIMENTAL_FEATURE, false,
+             "cub::ShuffleIndex",
+             "--use-experimental-features=non-uniform-groups");
     }
+    return;
   }
 }
 

clang/test/dpct/cub/warplevel/shuffle.cu

@@ -54,6 +54,9 @@ __global__ void ShuffleIndexKernel2(int* data) {
   // CHECK: output = dpct::experimental::select_from_sub_group(0xaaaaaaaa, item_ct1.get_sub_group(), input, 0);
   output = cub::ShuffleIndex<32>(input, 0, 0xaaaaaaaa);
   data[threadid] = output;
+  // CHECK: vec_output = dpct::select_from_sub_group(item_ct1.get_sub_group(), vec_input, 0);
+  float2 vec_input, vec_output;
+  vec_output = cub::ShuffleIndex<32>(vec_input, 0, 0xffffffff);
 }
 
 __global__ void ShuffleIndexKernel3(int* data) {

clang/test/dpct/cub/warplevel/shuffle_without_exp.cu

@@ -10,14 +10,18 @@ __global__ void ShuffleDownKernel(int *data) {
   int tid = cub::LaneId();
   unsigned mask = 0x8;
   int val = tid;
-  // CHECK: DPCT1007:{{.*}}: Migration of cub::ShuffleDown is not supported.
+// CHECK:   /*
+// CHECK:   DPCT1119:{{.*}}: Migration of cub::ShuffleDown is not supported, please try to remigrate with option: --use-experimental-features=non-uniform-groups.
+// CHECK:   */
   data[tid] = cub::ShuffleDown<8>(val, 3, 6, mask);
 }
 
 __global__ void ShuffleUpKernel(int *data) {
   int tid = cub::LaneId();
   unsigned mask = 0x8;
   int val = tid;
-  // CHECK: DPCT1007:{{.*}}: Migration of cub::ShuffleUp is not supported.
+// CHECK:   /*
+// CHECK:   DPCT1119:{{.*}}: Migration of cub::ShuffleUp is not supported, please try to remigrate with option: --use-experimental-features=non-uniform-groups.
+// CHECK:   */
   data[tid] = cub::ShuffleUp<8>(val, 3, 6, mask);
 }