[Scheduler] WIP (4)

Author: YWHyuk

PyTorchSimDevice/torch_openreg/csrc/Module.cpp

@@ -185,171 +185,6 @@ PyObject* _streamDestroy(PyObject* self, PyObject* arg) {
   END_HANDLE_TH_ERRORS
 }
 
-PyObject* _streamSynchronize(PyObject* self, PyObject* arg) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(THPUtils_checkLong(arg), "stream_synchronize expects an int");
-  orStream_t stream = reinterpret_cast<orStream_t>(THPUtils_unpackLong(arg));
-
-  orError_t err;
-  Py_BEGIN_ALLOW_THREADS
-  err = orStreamSynchronize(stream);
-  Py_END_ALLOW_THREADS
-
-  if (err != orSuccess) {
-    TORCH_CHECK(false, "Failed to synchronize stream");
-  }
-  Py_RETURN_NONE;
-  END_HANDLE_TH_ERRORS
-}
-
-PyObject* _streamQuery(PyObject* self, PyObject* arg) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(THPUtils_checkLong(arg), "stream_query expects an int");
-  orStream_t stream = reinterpret_cast<orStream_t>(THPUtils_unpackLong(arg));
-  orError_t err = orStreamQuery(stream);
-  if (err == orSuccess) {
-    Py_RETURN_TRUE;
-  } else {
-    Py_RETURN_FALSE;
-  }
-  END_HANDLE_TH_ERRORS
-}
-
-PyObject* _streamGetPriority(PyObject* self, PyObject* arg) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(THPUtils_checkLong(arg), "stream_get_priority expects an int");
-  orStream_t stream = reinterpret_cast<orStream_t>(THPUtils_unpackLong(arg));
-  int priority = 0;
-  orError_t err = orStreamGetPriority(stream, &priority);
-  if (err != orSuccess) {
-    TORCH_CHECK(false, "Failed to get stream priority");
-  }
-  return THPUtils_packInt32(priority);
-  END_HANDLE_TH_ERRORS
-}
-
-PyObject* _streamWaitEvent(PyObject* self, PyObject* args) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(PyTuple_Size(args) == 2, "stream_wait_event expects 2 arguments");
-  PyObject* stream_obj = PyTuple_GetItem(args, 0);
-  PyObject* event_obj = PyTuple_GetItem(args, 1);
-  TORCH_CHECK(THPUtils_checkLong(stream_obj), "stream must be an int");
-  TORCH_CHECK(THPUtils_checkLong(event_obj), "event must be an int");
-  orStream_t stream = reinterpret_cast<orStream_t>(THPUtils_unpackLong(stream_obj));
-  orEvent_t event = reinterpret_cast<orEvent_t>(THPUtils_unpackLong(event_obj));
-  orError_t err = orStreamWaitEvent(stream, event, 0);
-  if (err != orSuccess) {
-    TORCH_CHECK(false, "Failed to wait for event");
-  }
-  Py_RETURN_NONE;
-  END_HANDLE_TH_ERRORS
-}
-
-// Event functions
-PyObject* _eventCreate(PyObject* self, PyObject* noargs) {
-  HANDLE_TH_ERRORS
-  torch::utils::device_lazy_init(at::kPrivateUse1);
-  orEvent_t event = nullptr;
-  orError_t err = orEventCreate(&event);
-  if (err != orSuccess) {
-    TORCH_CHECK(false, "Failed to create event");
-  }
-  return THPUtils_packInt64(reinterpret_cast<int64_t>(event));
-  END_HANDLE_TH_ERRORS
-}
-
-PyObject* _eventCreateWithFlags(PyObject* self, PyObject* arg) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(THPUtils_checkLong(arg), "event_create_with_flags expects an int");
-  unsigned int flags = static_cast<unsigned int>(THPUtils_unpackLong(arg));
-
-  torch::utils::device_lazy_init(at::kPrivateUse1);
-  orEvent_t event = nullptr;
-  orError_t err = orEventCreateWithFlags(&event, flags);
-  if (err != orSuccess) {
-    TORCH_CHECK(false, "Failed to create event with flags");
-  }
-  return THPUtils_packInt64(reinterpret_cast<int64_t>(event));
-  END_HANDLE_TH_ERRORS
-}
-
-PyObject* _eventDestroy(PyObject* self, PyObject* arg) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(THPUtils_checkLong(arg), "event_destroy expects an int");
-  orEvent_t event = reinterpret_cast<orEvent_t>(THPUtils_unpackLong(arg));
-  orError_t err = orEventDestroy(event);
-  if (err != orSuccess) {
-    TORCH_CHECK(false, "Failed to destroy event");
-  }
-  Py_RETURN_NONE;
-  END_HANDLE_TH_ERRORS
-}
-
-PyObject* _eventRecord(PyObject* self, PyObject* args) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(PyTuple_Size(args) == 2, "event_record expects 2 arguments");
-  PyObject* event_obj = PyTuple_GetItem(args, 0);
-  PyObject* stream_obj = PyTuple_GetItem(args, 1);
-  TORCH_CHECK(THPUtils_checkLong(event_obj), "event must be an int");
-  TORCH_CHECK(THPUtils_checkLong(stream_obj), "stream must be an int");
-  orEvent_t event = reinterpret_cast<orEvent_t>(THPUtils_unpackLong(event_obj));
-  orStream_t stream = reinterpret_cast<orStream_t>(THPUtils_unpackLong(stream_obj));
-  orError_t err = orEventRecord(event, stream);
-  if (err != orSuccess) {
-    TORCH_CHECK(false, "Failed to record event");
-  }
-  Py_RETURN_NONE;
-  END_HANDLE_TH_ERRORS
-}
-
-PyObject* _eventSynchronize(PyObject* self, PyObject* arg) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(THPUtils_checkLong(arg), "event_synchronize expects an int");
-  orEvent_t event = reinterpret_cast<orEvent_t>(THPUtils_unpackLong(arg));
-
-  orError_t err;
-  Py_BEGIN_ALLOW_THREADS
-  err = orEventSynchronize(event);
-  Py_END_ALLOW_THREADS
-
-  if (err != orSuccess) {
-    TORCH_CHECK(false, "Failed to synchronize event");
-  }
-  Py_RETURN_NONE;
-  END_HANDLE_TH_ERRORS
-}
-
-PyObject* _eventQuery(PyObject* self, PyObject* arg) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(THPUtils_checkLong(arg), "event_query expects an int");
-  orEvent_t event = reinterpret_cast<orEvent_t>(THPUtils_unpackLong(arg));
-  orError_t err = orEventQuery(event);
-  if (err == orSuccess) {
-    Py_RETURN_TRUE;
-  } else {
-    Py_RETURN_FALSE;
-  }
-  END_HANDLE_TH_ERRORS
-}
-
-PyObject* _eventElapsedTime(PyObject* self, PyObject* args) {
-  HANDLE_TH_ERRORS
-  TORCH_CHECK(PyTuple_Size(args) == 2, "event_elapsed_time expects 2 arguments");
-  PyObject* start_obj = PyTuple_GetItem(args, 0);
-  PyObject* end_obj = PyTuple_GetItem(args, 1);
-  TORCH_CHECK(THPUtils_checkLong(start_obj), "start event must be an int");
-  TORCH_CHECK(THPUtils_checkLong(end_obj), "end event must be an int");
-  orEvent_t start = reinterpret_cast<orEvent_t>(THPUtils_unpackLong(start_obj));
-  orEvent_t end = reinterpret_cast<orEvent_t>(THPUtils_unpackLong(end_obj));
-  float ms = 0.0f;
-  orError_t err = orEventElapsedTime(&ms, start, end);
-  if (err != orSuccess) {
-    TORCH_CHECK(false, "Failed to get elapsed time");
-  }
-  return PyFloat_FromDouble(static_cast<double>(ms));
-  END_HANDLE_TH_ERRORS
-}
-
 PyObject* _deviceSynchronize(PyObject* self, PyObject* noargs) {
   HANDLE_TH_ERRORS
   torch::utils::device_lazy_init(at::kPrivateUse1);
@@ -421,20 +256,8 @@ static PyMethodDef methods[] = {
     {"get_amp_supported_dtype", _getAmpSupportedDtype, METH_NOARGS, nullptr},
     // Stream functions
     {"_stream_create", _streamCreate, METH_NOARGS, nullptr},
-    {"_stream_create_with_priority", _streamCreateWithPriority, METH_VARARGS, nullptr},
     {"_stream_destroy", _streamDestroy, METH_O, nullptr},
-    {"_stream_synchronize", _streamSynchronize, METH_O, nullptr},
-    {"_stream_query", _streamQuery, METH_O, nullptr},
-    {"_stream_get_priority", _streamGetPriority, METH_O, nullptr},
-    {"_stream_wait_event", _streamWaitEvent, METH_VARARGS, nullptr},
-    // Event functions
-    {"_event_create", _eventCreate, METH_NOARGS, nullptr},
-    {"_event_create_with_flags", _eventCreateWithFlags, METH_O, nullptr},
-    {"_event_destroy", _eventDestroy, METH_O, nullptr},
-    {"_event_record", _eventRecord, METH_VARARGS, nullptr},
-    {"_event_synchronize", _eventSynchronize, METH_O, nullptr},
-    {"_event_query", _eventQuery, METH_O, nullptr},
-    {"_event_elapsed_time", _eventElapsedTime, METH_VARARGS, nullptr},
+
     // Device functions
     {"_device_synchronize", _deviceSynchronize, METH_NOARGS, nullptr},
     // Stream task functions

PyTorchSimDevice/torch_openreg/openreg/__init__.py

@@ -8,6 +8,7 @@
 from .extension_device_interface import ExtensionDeviceInterface
 
 _initialized = False
+_default_streams = {}  # Dictionary to store default streams per device
 
 
 class device:
@@ -64,36 +65,21 @@ def _lazy_init():
     register_interface_for_device(custom_device(), ExtensionDeviceInterface)
     _initialized = True
 
+    # Create default streams for all devices
+    num_devices = device_count()
+    for device_idx in range(num_devices):
+        _default_streams[device_idx] = Stream()
 
 class Stream:
     """Wrapper for OpenReg stream."""
 
-    def __init__(self, priority=None, flags=0):
-        if priority is not None:
-            self._stream = torch_openreg._C._stream_create_with_priority(flags, priority)
-        else:
-            self._stream = torch_openreg._C._stream_create()
+    def __init__(self, flags=0):
+        self._stream = torch_openreg._C._stream_create()
 
     def __del__(self):
         if hasattr(self, '_stream'):
             torch_openreg._C._stream_destroy(self._stream)
 
-    def synchronize(self):
-        """Wait for all operations in the stream to complete."""
-        torch_openreg._C._stream_synchronize(self._stream)
-
-    def query(self):
-        """Check if all operations in the stream have completed."""
-        return torch_openreg._C._stream_query(self._stream)
-
-    def wait_event(self, event):
-        """Make this stream wait for an event."""
-        torch_openreg._C._stream_wait_event(self._stream, event._event)
-
-    def get_priority(self):
-        """Get the priority of the stream."""
-        return torch_openreg._C._stream_get_priority(self._stream)
-
     def launch_kernel(self, task):
         """Add a Python callable kernel to this stream.
 
@@ -108,74 +94,43 @@ def cdata(self):
         return self._stream
 
 
-class Event:
-    """Wrapper for OpenReg event."""
-
-    def __init__(self, enable_timing=False):
-        if enable_timing:
-            # orEventEnableTiming = 0x1
-            self._event = torch_openreg._C._event_create_with_flags(0x1)
-        else:
-            self._event = torch_openreg._C._event_create()
-
-    def __del__(self):
-        if hasattr(self, '_event'):
-            torch_openreg._C._event_destroy(self._event)
-
-    def record(self, stream=None):
-        """Record the event in a stream."""
-        if stream is None:
-            # Use default stream (stream 0)
-            stream = Stream()
-        torch_openreg._C._event_record(self._event, stream._stream)
-
-    def synchronize(self):
-        """Wait for the event to complete."""
-        torch_openreg._C._event_synchronize(self._event)
-
-    def query(self):
-        """Check if the event has completed."""
-        return torch_openreg._C._event_query(self._event)
-
-    def elapsed_time(self, start_event):
-        """Get the elapsed time between two events in milliseconds."""
-        return torch_openreg._C._event_elapsed_time(start_event._event, self._event)
-
-    @property
-    def cdata(self):
-        """Get the underlying event pointer (for internal use)."""
-        return self._event
-
-
 def synchronize():
     """Synchronize all streams on the current device."""
     torch_openreg._C._device_synchronize()
 
 
-def stream(priority=None, flags=0):
+def stream(flags=0):
     """Create a new stream.
 
     Args:
-        priority: Stream priority (optional)
         flags: Stream flags (optional)
 
     Returns:
         Stream: A new stream object
     """
-    return Stream(priority=priority, flags=flags)
+    return Stream(flags=flags)
 
+def default_stream(device=None):
+    _lazy_init()
+    if device is None:
+        device_idx = current_device()
+    else:
+        device_idx = torch.accelerator._get_device_index(device, optional=True)
+        if device_idx < 0:
+            device_idx = current_device()
 
-def event(enable_timing=False):
-    """Create a new event.
+    if device_idx not in _default_streams:
+        # Create default stream if it doesn't exist
+        _default_streams[device_idx] = Stream()
 
-    Args:
-        enable_timing: Whether to enable timing for the event
+    return _default_streams[device_idx]
 
-    Returns:
-        Event: A new event object
-    """
-    return Event(enable_timing=enable_timing)
 
+def launch_kernel(task, stream=None):
+    _lazy_init()
+    if stream is None:
+        stream = default_stream()
+    stream.launch_kernel(task)
 
 from .random import *  # noqa: F403
 from .amp import *
@@ -200,9 +155,7 @@ def event(enable_timing=False):
     "get_autocast_dtype",
     "set_autocast_dtype",
     "get_amp_supported_dtype",
-    "Stream",
-    "Event",
     "stream",
-    "event",
+    "launch_kernel",
     "synchronize",
 ]