Add flattened_field_count property to RangeModel and NdRangeModel.

Author: Diptorup Deb

numba_dpex/core/datamodel/models.py

@@ -20,6 +20,30 @@
 )
 
 
+def _get_flattened_member_count(ty):
+    """Return the number of fields in an instance of a given StructModel."""
+    flattened_member_count = 0
+    members = ty._members
+    for member in members:
+        if isinstance(member, types.UniTuple):
+            flattened_member_count += member.count
+        elif isinstance(
+            member,
+            (
+                types.scalars.Integer,
+                types.misc.PyObject,
+                types.misc.RawPointer,
+                types.misc.CPointer,
+                types.misc.MemInfoPointer,
+            ),
+        ):
+            flattened_member_count += 1
+        else:
+            raise UnreachableError
+
+    return flattened_member_count
+
+
 class GenericPointerModel(PrimitiveModel):
     def __init__(self, dmm, fe_type):
         adrsp = (
@@ -68,26 +92,7 @@ def __init__(self, dmm, fe_type):
     @property
     def flattened_field_count(self):
         """Return the number of fields in an instance of a USMArrayModel."""
-        flattened_member_count = 0
-        members = self._members
-        for member in members:
-            if isinstance(member, types.UniTuple):
-                flattened_member_count += member.count
-            elif isinstance(
-                member,
-                (
-                    types.scalars.Integer,
-                    types.misc.PyObject,
-                    types.misc.RawPointer,
-                    types.misc.CPointer,
-                    types.misc.MemInfoPointer,
-                ),
-            ):
-                flattened_member_count += 1
-            else:
-                raise UnreachableError
-
-        return flattened_member_count
+        return _get_flattened_member_count(self)
 
 
 class DpnpNdArrayModel(StructModel):
@@ -121,26 +126,7 @@ def __init__(self, dmm, fe_type):
     @property
     def flattened_field_count(self):
         """Return the number of fields in an instance of a DpnpNdArrayModel."""
-        flattened_member_count = 0
-        members = self._members
-        for member in members:
-            if isinstance(member, types.UniTuple):
-                flattened_member_count += member.count
-            elif isinstance(
-                member,
-                (
-                    types.scalars.Integer,
-                    types.misc.PyObject,
-                    types.misc.RawPointer,
-                    types.misc.CPointer,
-                    types.misc.MemInfoPointer,
-                ),
-            ):
-                flattened_member_count += 1
-            else:
-                raise UnreachableError
-
-        return flattened_member_count
+        return _get_flattened_member_count(self)
 
 
 class SyclQueueModel(StructModel):
@@ -211,6 +197,11 @@ def __init__(self, dmm, fe_type):
         ]
         super(RangeModel, self).__init__(dmm, fe_type, members)
 
+    @property
+    def flattened_field_count(self):
+        """Return the number of fields in an instance of a RangeModel."""
+        return _get_flattened_member_count(self)
+
 
 class NdRangeModel(StructModel):
     """The native data model for a
@@ -229,6 +220,11 @@ def __init__(self, dmm, fe_type):
         ]
         super(NdRangeModel, self).__init__(dmm, fe_type, members)
 
+    @property
+    def flattened_field_count(self):
+        """Return the number of fields in an instance of a NdRangeModel."""
+        return _get_flattened_member_count(self)
+
 
 def _init_data_model_manager() -> datamodel.DataModelManager:
     """Initializes a DpexKernelTarget-specific data model manager.