Add support for cuStreamGetCtx

oneapi-src · Jan 17, 2025 · f5dd74c · f5dd74c
1 parent be925b2
commit f5dd74c
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
diff --git a/clang/lib/DPCT/RulesLang/APINamesStream.inc b/clang/lib/DPCT/RulesLang/APINamesStream.inc
@@ -31,6 +31,12 @@ ASSIGNABLE_FACTORY(HEADER_INSERT_FACTORY(
                                    MEMBER_CALL(ARG(0), true, "wait"),
                                    CALL(ARG(1), ARG(0), ARG("0"), ARG(2)))))))
 
+ASSIGNABLE_FACTORY(FEATURE_REQUEST_FACTORY(
+    HelperFeatureEnum::device_ext,
+    ASSIGN_FACTORY_ENTRY("cuStreamGetCtx", DEREF(1),
+                         CALL(MapNames::getDpctNamespace() + "get_device_id",
+                            MEMBER_CALL(ARG(0), true, "get_device")))))
+
 ASSIGNABLE_FACTORY(CONDITIONAL_FACTORY_ENTRY(
     UseExtGraph,
     CALL_FACTORY_ENTRY("cudaStreamBeginCapture",

diff --git a/clang/lib/DPCT/RulesLang/RulesLang.cpp b/clang/lib/DPCT/RulesLang/RulesLang.cpp
@@ -4178,7 +4178,7 @@ void StreamAPICallRule::registerMatcher(MatchFinder &MF) {
         "cudaStreamIsCapturing", "cudaStreamQuery", "cudaStreamWaitEvent",
         "cudaStreamAddCallback", "cuStreamCreate", "cuStreamSynchronize",
         "cuStreamWaitEvent", "cuStreamDestroy_v2", "cuStreamAttachMemAsync",
-        "cuStreamAddCallback", "cuStreamQuery");
+        "cuStreamAddCallback", "cuStreamQuery", "cuStreamGetCtx");
   };
 
   MF.addMatcher(

diff --git a/clang/test/dpct/driver-stream-and-event.cu b/clang/test/dpct/driver-stream-and-event.cu
@@ -183,3 +183,9 @@ unsigned getEventFlags(bool enabledSyncBlock) {
 
   return flags;
 }
+
+// CHECK: void test_stream_and_context(dpct::queur_ptr stream, int& context);
+void test_stream_and_context(CUstream stream, CUcontext& context) {
+  // CHECK: context = dpct::get_device_id(stream->get_device());
+  cuStreamGetCtx(stream, &context);
+}