defines offloading helper types

Author: kaushikcfd

pyop2/offload_utils.py

@@ -0,0 +1,32 @@
+from enum import IntEnum
+
+
+class DataAvailability(IntEnum):
+    """
+    Indicates whether the device or host contains valid data.
+    """
+    AVAILABLE_ON_HOST_ONLY = 1
+    AVAILABLE_ON_DEVICE_ONLY = 2
+    AVAILABLE_ON_BOTH = 3
+
+
+class OffloadMixin:
+    def get_availability(self):
+        raise NotImplementedError
+
+    def ensure_availability_on_host(self):
+        raise NotImplementedError
+
+    def ensure_availaibility_on_device(self):
+        raise NotImplementedError
+
+    def is_available_on_host(self):
+        return bool(self.get_availability() & AVAILABLE_ON_HOST_ONLY)
+
+    def is_available_on_device(self):
+        return bool(self.get_availability() & AVAILABLE_ON_DEVICE_ONLY)
+
+
+AVAILABLE_ON_HOST_ONLY = DataAvailability.AVAILABLE_ON_HOST_ONLY
+AVAILABLE_ON_DEVICE_ONLY = DataAvailability.AVAILABLE_ON_DEVICE_ONLY
+AVAILABLE_ON_BOTH = DataAvailability.AVAILABLE_ON_BOTH