issue/571 - DeviceEvent implementation and interface

Author: wooway777

include/infinicore.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "infinicore/device_event.hpp"
 #include "infinicore/nn.hpp"
 #include "infinicore/ops.hpp"
 #include "infinicore/tensor.hpp"

include/infinicore/context/context.hpp

@@ -38,6 +38,7 @@ bool queryEvent(infinirtEvent_t event);
 void synchronizeEvent(infinirtEvent_t event);
 void destroyEvent(infinirtEvent_t event);
 float elapsedTime(infinirtEvent_t start, infinirtEvent_t end);
+void streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event);
 
 } // namespace context
 

include/infinicore/device_event.hpp

@@ -0,0 +1,125 @@
+#pragma once
+
+#include "device.hpp"
+#include "infinirt.h"
+#include <memory>
+#include <stdexcept>
+
+namespace infinicore {
+
+/**
+ * @brief A device event for timing operations and synchronization across devices.
+ *
+ * Similar to torch.cuda.Event, this class provides functionality to:
+ * - Record events on specific device streams
+ * - Synchronize with events
+ * - Measure elapsed time between events
+ * - Query event completion status
+ * - Make streams wait for events
+ */
+class DeviceEvent {
+private:
+    infinirtEvent_t event_; // Underlying event handle
+    Device device_;         // Device where this event was created
+    bool is_recorded_;      // Whether the event has been recorded
+
+public:
+    /**
+     * @brief Construct a new DeviceEvent on the current device.
+     */
+    DeviceEvent();
+
+    /**
+     * @brief Construct a new DeviceEvent on the current device with specific flags.
+     * @param flags Event creation flags (e.g., for timing, blocking sync)
+     */
+    explicit DeviceEvent(uint32_t flags);
+
+    /**
+     * @brief Construct a new DeviceEvent on a specific device.
+     * @param device Target device for this event
+     */
+    explicit DeviceEvent(Device device);
+
+    /**
+     * @brief Construct a new DeviceEvent on a specific device with flags.
+     * @param device Target device for this event
+     * @param flags Event creation flags
+     */
+    DeviceEvent(Device device, uint32_t flags);
+
+    // Disallow copying
+    DeviceEvent(const DeviceEvent &) = delete;
+    DeviceEvent &operator=(const DeviceEvent &) = delete;
+
+    /**
+     * @brief Move constructor.
+     */
+    DeviceEvent(DeviceEvent &&other) noexcept;
+
+    /**
+     * @brief Move assignment operator.
+     */
+    DeviceEvent &operator=(DeviceEvent &&other) noexcept;
+
+    /**
+     * @brief Destroy the DeviceEvent and release underlying resources.
+     */
+    ~DeviceEvent();
+
+    /**
+     * @brief Record the event on the current stream of its device.
+     */
+    void record();
+
+    /**
+     * @brief Record the event on a specific stream.
+     * @param stream Stream to record the event on
+     */
+    void record(infinirtStream_t stream);
+
+    /**
+     * @brief Wait for the event to complete (blocking).
+     */
+    void synchronize();
+
+    /**
+     * @brief Check if the event has been completed.
+     * @return true if completed, false otherwise
+     */
+    bool query() const;
+
+    /**
+     * @brief Calculate elapsed time between this event and another event (in milliseconds).
+     * @param other The other event to compare with
+     * @return Elapsed time in milliseconds
+     * @throws std::runtime_error if events are on different devices or not recorded
+     */
+    float elapsed_time(const DeviceEvent &other) const;
+
+    /**
+     * @brief Make a stream wait for this event to complete.
+     * @param stream Stream to make wait for this event (nullptr for current stream)
+     */
+    void wait(infinirtStream_t stream = nullptr) const;
+
+    /**
+     * @brief Get the device where this event was created.
+     * @return Device associated with this event
+     */
+    Device device() const { return device_; }
+
+    /**
+     * @brief Get the underlying event handle.
+     * @return Raw event handle
+     */
+    infinirtEvent_t get() const { return event_; }
+
+    /**
+     * @brief Check if the event has been recorded.
+     * @return true if recorded, false otherwise
+     */
+    bool is_recorded() const { return is_recorded_; }
+};
+
+} // namespace infinicore

python/infinicore/__init__.py

@@ -1,7 +1,18 @@
 import contextlib
 
 import infinicore.nn as nn
+
+# Import context functions
+from infinicore.context import (
+    get_device,
+    get_device_count,
+    get_stream,
+    set_device,
+    sync_device,
+    sync_stream,
+)
 from infinicore.device import device
+from infinicore.device_event import DeviceEvent
 from infinicore.dtype import (
     bfloat16,
     bool,
@@ -47,8 +58,16 @@
     "nn",
     # Classes.
     "device",
+    "DeviceEvent",
     "dtype",
     "Tensor",
+    # Context functions.
+    "get_device",
+    "get_device_count",
+    "get_stream",
+    "set_device",
+    "sync_device",
+    "sync_stream",
     # Data Types.
     "bfloat16",
     "bool",

python/infinicore/context.py

@@ -0,0 +1,50 @@
+from infinicore.lib import _infinicore
+
+
+def get_device():
+    """Get the current active device.
+
+    Returns:
+        device: The current active device object
+    """
+    return _infinicore.get_device()
+
+
+def get_device_count(device_type):
+    """Get the number of available devices of a specific type.
+
+    Args:
+        device_type (str): The type of device to count (e.g., "cuda", "cpu", "npu")
+
+    Returns:
+        int: The number of available devices of the specified type
+    """
+    return _infinicore.get_device_count(device_type)
+
+
+def set_device(device):
+    """Set the current active device.
+
+    Args:
+        device: The device to set as active
+    """
+    _infinicore.set_device(device._underlying)
+
+
+def sync_stream():
+    """Synchronize the current stream."""
+    _infinicore.sync_stream()
+
+
+def sync_device():
+    """Synchronize the current device."""
+    _infinicore.sync_device()
+
+
+def get_stream():
+    """Get the current stream.
+
+    Returns:
+        stream: The current stream object
+    """
+    return _infinicore.get_stream()

python/infinicore/device_event.py

@@ -0,0 +1,95 @@
+import infinicore.device
+from infinicore.lib import _infinicore
+
+
+class DeviceEvent:
+    """A device event for timing operations and synchronization across devices.
+
+    Similar to torch.cuda.Event, this class provides functionality to:
+    - Record events on specific device streams
+    - Synchronize with events
+    - Measure elapsed time between events
+    - Query event completion status
+    - Make streams wait for events
+
+    Args:
+        device: Target device for this event. If None, uses current device.
+        flags: Event creation flags (e.g., for timing, blocking sync). Default is 0.
+        enable_timing: Whether the event should be created with timing enabled.
+    """
+
+    def __init__(self, device=None, enable_timing=True, flags=0):
+        if not enable_timing:
+            # You might want to handle this differently based on your flag system
+            flags = flags  # Adjust flags if timing is disabled
+
+        if device is None:
+            # Use current device
+            if flags == 0:
+                self._underlying = _infinicore.DeviceEvent()
+            else:
+                self._underlying = _infinicore.DeviceEvent(flags)
+        elif flags == 0:
+            # Construct with device only
+            self._underlying = _infinicore.DeviceEvent(device._underlying)
+        else:
+            # Construct with both device and flags
+            self._underlying = _infinicore.DeviceEvent(device._underlying, flags)
+
+    def record(self, stream=None):
+        """Record the event.
+
+        Args:
+            stream: Stream to record the event on. If None, uses current stream.
+        """
+        if stream is None:
+            self._underlying.record()
+        else:
+            self._underlying.record(stream)
+
+    def synchronize(self):
+        """Wait for the event to complete (blocking)."""
+        self._underlying.synchronize()
+
+    def query(self):
+        """Check if the event has been completed.
+
+        Returns:
+            bool: True if completed, False otherwise.
+        """
+        return self._underlying.query()
+
+    def elapsed_time(self, other):
+        """Calculate elapsed time between this event and another event.
+
+        Args:
+            other: The other DeviceEvent to compare with
+
+        Returns:
+            float: Elapsed time in milliseconds between this event and the other event
+
+        Raises:
+            RuntimeError: If events are on different devices or not recorded
+        """
+        return self._underlying.elapsed_time(other._underlying)
+
+    def wait(self, stream=None):
+        """Make a stream wait for this event to complete.
+
+        Args:
+            stream: Stream to make wait for this event. If None, uses current stream.
+        """
+        self._underlying.wait(stream)
+
+    @property
+    def device(self):
+        """Get the device where this event was created."""
+        return infinicore.device._from_infinicore_device(self._underlying.device)
+
+    @property
+    def is_recorded(self):
+        """Check if the event has been recorded."""
+        return self._underlying.is_recorded
+
+    def __repr__(self):
+        return f"DeviceEvent(device={self.device}, recorded={self.is_recorded})"

src/infinicore/context/context_impl.cc

@@ -58,7 +58,7 @@ ContextImpl &ContextImpl::singleton() {
 }
 
 ContextImpl::ContextImpl() {
-    std::vector<int> device_counter(size_t(Device::Type::COUNT));
+    std::vector<int> device_counter(static_cast<size_t>(Device::Type::COUNT));
     INFINICORE_CHECK_ERROR(infinirtGetAllDeviceCount(device_counter.data()));
 
     // Reserve runtime slot for all devices.
@@ -168,6 +168,10 @@ float elapsedTime(infinirtEvent_t start, infinirtEvent_t end) {
     return ContextImpl::singleton().getCurrentRuntime()->elapsedTime(start, end);
 }
 
+void streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) {
+    ContextImpl::singleton().getCurrentRuntime()->streamWaitEvent(stream, event);
+}
+
 } // namespace context
 
 } // namespace infinicore

src/infinicore/context/runtime/runtime.cc

@@ -128,6 +128,14 @@ float Runtime::elapsedTime(infinirtEvent_t start, infinirtEvent_t end) {
     return ms;
 }
 
+void Runtime::streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) {
+    // Use current stream if no specific stream is provided
+    if (stream == nullptr) {
+        stream = stream_;
+    }
+    INFINICORE_CHECK_ERROR(infinirtStreamWaitEvent(stream, event));
+}
+
 std::string Runtime::toString() const {
     return fmt::format("Runtime({})", device_.toString());
 }

src/infinicore/context/runtime/runtime.hpp

@@ -46,6 +46,7 @@ class Runtime {
     void synchronizeEvent(infinirtEvent_t event);
     void destroyEvent(infinirtEvent_t event);
     float elapsedTime(infinirtEvent_t start, infinirtEvent_t end);
+    void streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event);
 
     std::string toString() const;