[SYCLomatic] Migrate cudaMallocAsync and cudaFreeAsync

clang/include/clang/DPCT/DPCTOptions.inc

@@ -811,11 +811,12 @@ DPCT_ENUM_OPTION(
             "be accessed within a kernel using syntax similar to C++ global "
             "variables.\n",
             false),
-        DPCT_OPTION_ENUM_VALUE(
-            "virtual_mem", int(ExperimentalFeatures::Exp_VirtualMemory),
-            "Experimental extension that allows for mapping of an address range onto "
-            "multiple allocations of physical memory.",
-            false),
+        DPCT_OPTION_ENUM_VALUE("virtual_mem",
+                               int(ExperimentalFeatures::Exp_VirtualMemory),
+                               "Experimental extension that allows for mapping "
+                               "of an address range onto "
+                               "multiple allocations of physical memory.",
+                               false),
         DPCT_OPTION_ENUM_VALUE(
             "in_order_queue_events",
             int(ExperimentalFeatures::Exp_InOrderQueueEvents),
@@ -838,7 +839,13 @@ DPCT_ENUM_OPTION(
             "level_zero", int(ExperimentalFeatures::Exp_LevelZero),
             "Experimental migration feature that enables the use of Level Zero "
             "APIs to migrate target code, like CUDA Inter-Process "
-            "Communication (IPC) APIs.\n", false),
+            "Communication (IPC) APIs.\n",
+            false),
+        DPCT_OPTION_ENUM_VALUE("async_alloc",
+                               int(ExperimentalFeatures::Exp_AsyncAlloc),
+                               "Experimental extension that allows use of SYCL "
+                               "async allocation APIs.\n",
+                               false),
         DPCT_OPTION_ENUM_VALUE(
             "all", int(ExperimentalFeatures::Exp_All),
             "Enable all experimental extensions listed in this option.\n",

clang/lib/DPCT/AnalysisInfo.h

@@ -1352,6 +1352,9 @@ class DpctGlobalInfo {
   static bool useExtLevelZero() {
     return getUsingExperimental<ExperimentalFeatures::Exp_LevelZero>();
   }
+  static bool useExtAsyncAlloc() {
+    return getUsingExperimental<ExperimentalFeatures::Exp_AsyncAlloc>();
+  }
   static bool useExtPrefetch() {
     return getUsingExperimental<ExperimentalFeatures::Exp_Prefetch>();
   }

clang/lib/DPCT/CommandOption/ValidateArguments.h

@@ -103,6 +103,7 @@ enum class ExperimentalFeatures : unsigned int {
   Exp_NonStandardSYCLBuiltins,
   Exp_Prefetch,
   Exp_LevelZero,
+  Exp_AsyncAlloc,
   Exp_All
 };
 enum class HelperFuncPreference : unsigned int { NoQueueDevice = 0 };

clang/lib/DPCT/RuleInfra/CallExprRewriterCommon.h

@@ -1827,6 +1827,10 @@ inline auto UseExtLevelZero = [](const CallExpr *C) -> bool {
   return DpctGlobalInfo::useExtLevelZero();
 };
 
+inline auto UseExtAsyncAlloc = [](const CallExpr *C) -> bool {
+  return DpctGlobalInfo::useExtAsyncAlloc();
+};
+
 inline auto UseExtGraph = [](const CallExpr *C) -> bool {
   return DpctGlobalInfo::useExtGraph();
 };

clang/lib/DPCT/RulesInclude/HeaderTypes.inc

@@ -52,6 +52,8 @@ STD_HEADER(DL, "<dlfcn.h>")
 #endif
 STD_HEADER(SHMEM, "<ishmem.h>")
 STD_HEADER(SHMEMX, "<ishmemx.h>")
+STD_HEADER(AsyncAlloc,
+           "<sycl/ext/oneapi/experimental/async_alloc/async_alloc.hpp>")
 
 ONEDPL_HEADER(Algorithm, "<oneapi/dpl/algorithm>")
 ONEDPL_HEADER(Execution, "<oneapi/dpl/execution>")

clang/lib/DPCT/RulesLang/APINamesMemory.inc

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "RuleInfra/CallExprRewriterCommon.h"
 ASSIGNABLE_FACTORY(CONDITIONAL_FACTORY_ENTRY(
     checkIsUSM(),
     CONDITIONAL_FACTORY_ENTRY(
@@ -993,6 +994,42 @@ CONDITIONAL_FACTORY_ENTRY(
                                    "memcpy", false),
                                MEM_ARG(0), MEM_ARG(1), ARG(2), ARG(3))))))
 
+CONDITIONAL_FACTORY_ENTRY(
+    checkIsUSM(),
+    CONDITIONAL_FACTORY_ENTRY(
+        UseExtAsyncAlloc,
+        ASSIGNABLE_FACTORY(HEADER_INSERT_FACTORY(
+            HeaderType::HT_AsyncAlloc,
+            ASSIGN_FACTORY_ENTRY(
+                "cudaMallocAsync", DEREF(0),
+                CALL(MapNames::getClNamespace() +
+                         "ext::oneapi::experimental::async_malloc",
+                     DEREF(2),
+                     ARG(MapNames::getClNamespace() + "usm::alloc::device"),
+                     ARG(1))))),
+        UNSUPPORT_FACTORY_ENTRY(
+            "cudaMallocAsync", Diagnostics::TRY_EXPERIMENTAL_FEATURE,
+            ARG("cudaMallocAsync"),
+            ARG("--use-experimental-features=virtual_mem"))),
+    UNSUPPORT_FACTORY_ENTRY("cudaMallocAsync", Diagnostics::API_NOT_MIGRATED,
+                            ARG("cudaMallocAsync")))
+CONDITIONAL_FACTORY_ENTRY(
+    checkIsUSM(),
+    CONDITIONAL_FACTORY_ENTRY(
+        UseExtAsyncAlloc,
+        ASSIGNABLE_FACTORY(HEADER_INSERT_FACTORY(
+            HeaderType::HT_AsyncAlloc,
+            CALL_FACTORY_ENTRY("cudaFreeAsync",
+                               CALL(MapNames::getClNamespace() +
+                                        "ext::oneapi::experimental::async_free",
+                                    DEREF(1), ARG(0))))),
+        UNSUPPORT_FACTORY_ENTRY(
+            "cudaFreeAsync", Diagnostics::TRY_EXPERIMENTAL_FEATURE,
+            ARG("cudaFreeAsync"),
+            ARG("--use-experimental-features=virtual_mem"))),
+    UNSUPPORT_FACTORY_ENTRY("cudaFreeAsync", Diagnostics::API_NOT_MIGRATED,
+                            ARG("cudaFreeAsync")))
+
 #define CUDA_FREE(NAME)                                                        \
   CONDITIONAL_FACTORY_ENTRY(                                                   \
       hasManagedAttr(0),                                                       \

clang/lib/DPCT/RulesLang/RulesLang.cpp

@@ -5784,7 +5784,7 @@ void MemoryMigrationRule::mallocMigration(
   } else if (Name == "cudaHostAlloc" || Name == "cudaMallocHost" ||
              Name == "cuMemHostAlloc" || Name == "cuMemAllocHost_v2" ||
              Name == "cuMemAllocPitch_v2" || Name == "cudaMallocPitch" ||
-             Name == "cudaMallocMipmappedArray") {
+             Name == "cudaMallocMipmappedArray" || Name == "cudaMallocAsync") {
     ExprAnalysis EA(C);
     emplaceTransformation(EA.getReplacement());
     EA.applyAllSubExprRepl();
@@ -6815,7 +6815,7 @@ void MemoryMigrationRule::registerMatcher(MatchFinder &MF) {
         "cuMemsetD8_v2", "cuMemsetD8Async", "cudaMallocMipmappedArray",
         "cudaGetMipmappedArrayLevel", "cudaFreeMipmappedArray",
         "cudaMemcpyPeer", "cudaMemcpyPeerAsync", "cuMemcpyPeer",
-        "cuMemcpyPeerAsync");
+        "cuMemcpyPeerAsync", "cudaMallocAsync", "cudaFreeAsync");
   };
 
   MF.addMatcher(callExpr(allOf(callee(functionDecl(memoryAPI())), parentStmt()))
@@ -6905,7 +6905,8 @@ void MemoryMigrationRule::runRule(const MatchFinder::MatchResult &Result) {
         Name.compare("cudaMallocMipmappedArray") &&
         Name.compare("cudaGetMipmappedArrayLevel") &&
         Name.compare("cudaFreeMipmappedArray") && Name.compare("cudaMemcpy") &&
-        Name.compare("cudaFree") && Name.compare("cublasFree")) {
+        Name.compare("cudaFree") && Name.compare("cublasFree") &&
+        Name.compare("cudaMallocAsync") && Name.compare("cudaFreeAsync")) {
       requestFeature(HelperFeatureEnum::device_ext);
       insertAroundStmt(C, MapNames::getCheckErrorMacroName() + "(", ")");
     } else if (IsAssigned && !Name.compare("cudaMemAdvise") &&
@@ -6968,6 +6969,7 @@ MemoryMigrationRule::MemoryMigrationRule() {
                          const CallExpr *, const UnresolvedLookupExpr *, bool)>>
       Dispatcher{
           {"cudaMalloc", &MemoryMigrationRule::mallocMigration},
+          {"cudaMallocAsync", &MemoryMigrationRule::mallocMigration},
           {"cuMemAlloc_v2", &MemoryMigrationRule::mallocMigration},
           {"cudaHostAlloc", &MemoryMigrationRule::mallocMigration},
           {"cudaMallocHost", &MemoryMigrationRule::mallocMigration},
@@ -7030,6 +7032,7 @@ MemoryMigrationRule::MemoryMigrationRule() {
           {"cuMemcpyDtoA_v2", &MemoryMigrationRule::arrayMigration},
           {"cuMemcpyAtoA_v2", &MemoryMigrationRule::arrayMigration},
           {"cudaFree", &MemoryMigrationRule::freeMigration},
+          {"cudaFreeAsync", &MemoryMigrationRule::freeMigration},
           {"cuMemFree_v2", &MemoryMigrationRule::freeMigration},
           {"cudaFreeArray", &MemoryMigrationRule::freeMigration},
           {"cudaFreeMipmappedArray", &MemoryMigrationRule::freeMigration},

clang/lib/DPCT/SrcAPI/APINames.inc

@@ -237,8 +237,8 @@ ENTRY(cudaMemcpyToArray, cudaMemcpyToArray, true, NO_FLAG, P0, "Successful")
 ENTRY(cudaMemcpyToArrayAsync, cudaMemcpyToArrayAsync, true, NO_FLAG, P0, "Successful")
 
 // stream ordered memory allocator functions of runtime API
-ENTRY(cudaFreeAsync, cudaFreeAsync, false, NO_FLAG, P7, "comment")
-ENTRY(cudaMallocAsync, cudaMallocAsync, false, NO_FLAG, P7, "comment")
+ENTRY(cudaFreeAsync, cudaFreeAsync, true, NO_FLAG, P7, "comment")
+ENTRY(cudaMallocAsync, cudaMallocAsync, true, NO_FLAG, P7, "comment")
 ENTRY(cudaMallocFromPoolAsync, cudaMallocFromPoolAsync, false, NO_FLAG, P7, "comment")
 ENTRY(cudaMemPoolCreate, cudaMemPoolCreate, false, NO_FLAG, P4, "comment")
 ENTRY(cudaMemPoolDestroy, cudaMemPoolDestroy, false, NO_FLAG, P7, "comment")