Skip to content

Commit 185a61b

Browse files
author
Diptorup Deb
committed
Introduce and experimental data model manager.
1 parent 62b942d commit 185a61b

File tree

3 files changed

+41
-12
lines changed

3 files changed

+41
-12
lines changed

numba_dpex/experimental/launcher.py

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
from llvmlite import ir as llvmir
1313
from numba.core import cgutils, cpu, types
14+
from numba.core.datamodel import default_manager as numba_default_dmm
1415
from numba.extending import intrinsic, overload
1516

1617
from numba_dpex import config, dpjit
@@ -192,23 +193,27 @@ def create_llvm_values_for_index_space(
192193
ndim = indexer_argty.ndim
193194
grange_extents = []
194195
lrange_extents = []
195-
datamodel = self._kernel_targetctx.data_model_manager.lookup(
196-
indexer_argty
197-
)
196+
indexer_datamodel = numba_default_dmm.lookup(indexer_argty)
198197

199198
if isinstance(indexer_argty, RangeType):
200199
for dim_num in range(ndim):
201-
dim_pos = datamodel.get_field_position("dim" + str(dim_num))
200+
dim_pos = indexer_datamodel.get_field_position(
201+
"dim" + str(dim_num)
202+
)
202203
grange_extents.append(
203204
self._builder.extract_value(index_space_arg, dim_pos)
204205
)
205206
elif isinstance(indexer_argty, NdRangeType):
206207
for dim_num in range(ndim):
207-
gdim_pos = datamodel.get_field_position("gdim" + str(dim_num))
208+
gdim_pos = indexer_datamodel.get_field_position(
209+
"gdim" + str(dim_num)
210+
)
208211
grange_extents.append(
209212
self._builder.extract_value(index_space_arg, gdim_pos)
210213
)
211-
ldim_pos = datamodel.get_field_position("ldim" + str(dim_num))
214+
ldim_pos = indexer_datamodel.get_field_position(
215+
"ldim" + str(dim_num)
216+
)
212217
lrange_extents.append(
213218
self._builder.extract_value(index_space_arg, ldim_pos)
214219
)

numba_dpex/experimental/models.py

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,38 @@
22
#
33
# SPDX-License-Identifier: Apache-2.0
44

5-
"""Provides Numba datamodel for the numba_dpex types introduced in the
5+
"""Provides the Numba data models for the numba_dpex types introduced in the
66
numba_dpex.experimental module.
77
"""
88

9-
from numba.core.datamodel import models
9+
from numba.core.datamodel import DataModelManager, models
1010
from numba.core.extending import register_model
1111

12-
from numba_dpex.core.datamodel.models import dpex_data_model_manager as dmm
12+
import numba_dpex.core.datamodel.models as dpex_core_models
1313

1414
from .types import KernelDispatcherType
1515

16-
# Register the types and datamodel in the DpexKernelTargetContext
17-
dmm.register(KernelDispatcherType, models.OpaqueModel)
1816

19-
# Register the types and datamodel in the DpexTargetContext
17+
def _init_exp_data_model_manager() -> DataModelManager:
18+
"""Initializes a DpexExpKernelTarget-specific data model manager.
19+
20+
Extends the DpexKernelTargetContext's datamodel manager with all
21+
experimental types that are getting added to the kernel API.
22+
23+
Returns:
24+
DataModelManager: A numba-dpex DpexExpKernelTarget-specific data model
25+
manager
26+
"""
27+
28+
dmm = dpex_core_models.dpex_data_model_manager.copy()
29+
30+
# Register the types and data model in the DpexExpTargetContext
31+
# Add here...
32+
33+
return dmm
34+
35+
36+
exp_dmm = _init_exp_data_model_manager()
37+
38+
# Register any new type that should go into numba.core.datamodel.default_manager
2039
register_model(KernelDispatcherType)(models.OpaqueModel)

numba_dpex/experimental/target.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
DpexKernelTargetContext,
1717
DpexKernelTypingContext,
1818
)
19+
from numba_dpex.experimental.models import exp_dmm
1920

2021

2122
# pylint: disable=R0903
@@ -51,6 +52,10 @@ class DpexExpKernelTargetContext(DpexKernelTargetContext):
5152
they are stable enough to be migrated to DpexKernelTargetContext.
5253
"""
5354

55+
def __init__(self, typingctx, target=DPEX_KERNEL_EXP_TARGET_NAME):
56+
super().__init__(typingctx, target)
57+
self.data_model_manager = exp_dmm
58+
5459

5560
class DpexExpKernelTarget(TargetDescriptor):
5661
"""

0 commit comments

Comments
 (0)