diff --git a/loopy/kernel/__init__.py b/loopy/kernel/__init__.py index d618d4b0dd9adc0ccb50827dce914571538a62a4..5701bf77254277501e86e32cdef0e01edaa90b29 100644 --- a/loopy/kernel/__init__.py +++ b/loopy/kernel/__init__.py @@ -259,7 +259,7 @@ class LoopKernel(RecordWithoutPickling): # }}} from loopy.types import to_loopy_type - index_dtype = to_loopy_type(index_dtype).with_target(target) + index_dtype = to_loopy_type(index_dtype, target=target) if not index_dtype.is_integral(): raise TypeError("index_dtype must be an integer") if np.iinfo(index_dtype.numpy_dtype).min >= 0: diff --git a/loopy/kernel/array.py b/loopy/kernel/array.py index 99bbc7bf9e782fe7995c20d8d3482602c9874dc9..11392f255daa22126b519a0b9394917398048f6e 100644 --- a/loopy/kernel/array.py +++ b/loopy/kernel/array.py @@ -576,6 +576,7 @@ class ArrayBase(Record): def __init__(self, name, dtype=None, shape=None, dim_tags=None, offset=0, dim_names=None, strides=None, order=None, for_atomic=False, + target=None, **kwargs): """ All of the following are optional. Specify either strides or shape. @@ -659,7 +660,7 @@ class ArrayBase(Record): from loopy.types import to_loopy_type dtype = to_loopy_type(dtype, allow_auto=True, allow_none=True, - for_atomic=for_atomic) + for_atomic=for_atomic, target=target) strides_known = strides is not None and strides is not lp.auto shape_known = shape is not None and shape is not lp.auto diff --git a/loopy/kernel/creation.py b/loopy/kernel/creation.py index ff3bf16bcf32b26b1865d350aefbef80ec4e4554..e224bd0442233f9321293b4f53f39047507150b3 100644 --- a/loopy/kernel/creation.py +++ b/loopy/kernel/creation.py @@ -1244,7 +1244,8 @@ def create_temporaries(knl, default_order): scope=lp.auto, base_indices=lp.auto, shape=lp.auto, - order=default_order) + order=default_order, + target=knl.target) if isinstance(insn, Assignment): insn = insn.copy(temp_var_type=None) diff --git a/loopy/kernel/data.py b/loopy/kernel/data.py index 004fae7f9664ff62c34a994671ea792e4eddc836..463d10b4c9e06209f26261bb3ddf328520390e31 100644 --- a/loopy/kernel/data.py +++ b/loopy/kernel/data.py @@ -203,10 +203,12 @@ class KernelArgument(Record): def __init__(self, **kwargs): kwargs["name"] = intern(kwargs.pop("name")) + target = kwargs.pop("target", None) + dtype = kwargs.pop("dtype", None) from loopy.types import to_loopy_type kwargs["dtype"] = to_loopy_type( - dtype, allow_auto=True, allow_none=True) + dtype, allow_auto=True, allow_none=True, target=target) Record.__init__(self, **kwargs) @@ -243,11 +245,11 @@ class ImageArg(ArrayBase, KernelArgument): class ValueArg(KernelArgument): - def __init__(self, name, dtype=None, approximately=1000): - from loopy.types import to_loopy_type + def __init__(self, name, dtype=None, approximately=1000, target=None): KernelArgument.__init__(self, name=name, - dtype=to_loopy_type(dtype, allow_auto=True, allow_none=True), - approximately=approximately) + dtype=dtype, + approximately=approximately, + target=target) def __str__(self): import loopy as lp diff --git a/loopy/kernel/tools.py b/loopy/kernel/tools.py index 7e9bd549fede6abf6d4d5db99896063b34246793..e005d5ab2b10040ff2ea1601492b8347147b9451 100644 --- a/loopy/kernel/tools.py +++ b/loopy/kernel/tools.py @@ -70,7 +70,7 @@ def _add_dtypes(knl, dtype_dict): for arg in knl.args: new_dtype = dtype_dict.pop(arg.name, None) if new_dtype is not None: - new_dtype = to_loopy_type(new_dtype) + new_dtype = to_loopy_type(new_dtype, target=knl.target) if arg.dtype is not None and arg.dtype != new_dtype: raise RuntimeError( "argument '%s' already has a different dtype " diff --git a/loopy/types.py b/loopy/types.py index b897d9f700b198e73d95a09c7d459ed2d7f877b1..f095d1d58f9eaebb7dcc9c8d41afa73951f2ba84 100644 --- a/loopy/types.py +++ b/loopy/types.py @@ -177,7 +177,8 @@ class AtomicNumpyType(NumpyType, AtomicType): # }}} -def to_loopy_type(dtype, allow_none=False, allow_auto=False, for_atomic=False): +def to_loopy_type(dtype, allow_none=False, allow_auto=False, for_atomic=False, + target=None): from loopy.kernel.data import auto if allow_none and dtype is None: return dtype @@ -192,10 +193,13 @@ def to_loopy_type(dtype, allow_none=False, allow_auto=False, for_atomic=False): except Exception: pass + if numpy_dtype is None and target is not None and isinstance(dtype, str): + numpy_dtype = target.get_dtype_registry().get_or_register_dtype(dtype) + if isinstance(dtype, LoopyType): if for_atomic: if isinstance(dtype, NumpyType): - return AtomicNumpyType(dtype.dtype) + return AtomicNumpyType(dtype.dtype, target=target) elif not isinstance(dtype, AtomicType): raise LoopyError("do not know how to convert '%s' to an atomic type" % dtype) @@ -204,9 +208,9 @@ def to_loopy_type(dtype, allow_none=False, allow_auto=False, for_atomic=False): elif numpy_dtype is not None: if for_atomic: - return AtomicNumpyType(numpy_dtype) + return AtomicNumpyType(numpy_dtype, target=target) else: - return NumpyType(numpy_dtype) + return NumpyType(numpy_dtype, target=target) else: raise TypeError("dtype must be a LoopyType, or convertible to one, "