Commit d2f1f981 authored by Pierre Augier's avatar Pierre Augier
Browse files

mem_layout for arrays: C_or_F (default), C, F and strided

parent 57cd941a123a
......@@ -26,7 +26,7 @@ import inspect
from transonic.analyses.extast import unparse, ast, FunctionDef, Name
from transonic.signatures import make_signatures_from_typehinted_func
from transonic.typing import format_type_as_backend_type
from transonic.typing import format_type_as_backend_type, MemLayout
from .base import BackendAOT, TypeHintRemover, format_str
from .base_jit import SubBackendJIT
......@@ -49,15 +49,15 @@ class TypeFormatterCython(TypeFormatter):
return f"cython.{name}"
return name
def make_array_code(self, dtype, ndim, memview):
def make_array_code(self, dtype, ndim, memview, mem_layout):
dtype = normalize_type_name_for_array(dtype.__name__)
if ndim == 0:
return dtype
if memview:
return memoryview_type(dtype, ndim)
return memoryview_type(dtype, ndim, mem_layout)
else:
return np_ndarray_type(dtype, ndim)
return np_ndarray_type(dtype, ndim, mem_layout)
def make_dict_code(self, type_keys, type_values, **kwargs):
return "dict"
......@@ -72,13 +72,27 @@ class TypeFormatterCython(TypeFormatter):
return "tuple"
def memoryview_type(dtype, ndim) -> str:
end = "[" + ", ".join(":" * ndim) + "]"
return f"{dtype}_t{end}"
def np_ndarray_type(dtype, ndim) -> str:
return f"np.ndarray[{dtype}_t, ndim={ndim}]"
def memoryview_type(dtype, ndim, mem_layout) -> str:
ndim_F = 0
ndim_C = 0
if mem_layout is MemLayout.C:
ndim_C = 1
ndim -= 1
elif mem_layout is MemLayout.F:
ndim_F = 1
ndim -= 1
end = ", ".join(["::1"] * ndim_F + [":"] * ndim + ["::1"] * ndim_C)
return f"{dtype}_t[{end}]"
def np_ndarray_type(dtype, ndim, mem_layout) -> str:
if mem_layout is MemLayout.C:
mode = ', mode="c"'
elif mem_layout is MemLayout.F:
mode = ', mode="f"'
else:
mode = ""
return f"np.ndarray[{dtype}_t, ndim={ndim}{mode}]"
class HeaderFunction:
......
import numpy as np
from transonic import Array
from transonic.backends import backends
backend = backends["cython"]
type_formatter = backend.type_formatter
def compare(dtype, ndim, memview, mem_layout, result):
A = Array[dtype, ndim, memview, mem_layout]
assert A.format_as_backend_type(type_formatter) == result
def test_memview():
memview = "memview"
compare(int, "2d", memview, "C", "np.int_t[:, ::1]")
compare(int, "3d", memview, "strided", "np.int_t[:, :, :]")
compare(np.int32, "2d", memview, "F", "np.int32_t[::1, :]")
def test_array():
memview = None
compare(int, "2d", memview, "C", 'np.ndarray[np.int_t, ndim=2, mode="c"]')
compare(int, "3d", memview, "strided", "np.ndarray[np.int_t, ndim=3]")
compare(
np.int32, "2d", memview, "F", 'np.ndarray[np.int32_t, ndim=2, mode="f"]'
)
import numpy as np
from transonic import Array
from .typing import base_type_formatter
def compare(dtype, ndim, mem_layout, result):
A = Array[dtype, ndim, mem_layout]
assert A.format_as_backend_type(base_type_formatter) == result
def test_array():
compare(int, "2d", "C", "int[:, :] order(C)")
compare(int, "3d", "strided", "int[::, ::, ::]")
compare(np.int32, "2d", "F", "int32[:, :] order(F)")
from transonic.typing import format_type_as_backend_type
from transonic.typing import format_type_as_backend_type, MemLayout
normalized_types = {"float": "float64", "complex": "complex128"}
......@@ -13,11 +13,19 @@ class TypeFormatter:
except KeyError:
return name
def make_array_code(self, dtype, ndim, memview):
base = self.normalize_type_name(dtype.__name__)
def make_array_code(self, dtype, ndim, memview, mem_layout):
dtype = self.normalize_type_name(dtype.__name__)
if ndim == 0:
return base
return base + f"[{', '.join(':'*ndim)}]"
return dtype
one_dim = ":"
if mem_layout is MemLayout.strided:
one_dim = ["::"]
result = f"{dtype}[{', '.join(one_dim*ndim)}]"
if mem_layout is MemLayout.C:
result += " order(C)"
elif mem_layout is MemLayout.F:
result += " order(F)"
return result
def make_dict_code(self, type_keys, type_values, **kwargs):
key = format_type_as_backend_type(type_keys, self, **kwargs)
......
......@@ -24,8 +24,7 @@ def func0(a, b):
func0_boosted = boost(func0)
A = Union[int, Array[int, "1d"]]
A = Union[int, Array[int, "1d", "C"]]
@boost
......
......@@ -55,6 +55,7 @@ Internal API
"""
import re
from enum import Enum, auto
import numpy as np
......@@ -202,7 +203,18 @@ class UnionVar(TemplateVar):
class Meta(type):
pass
def __call__(cls, *args, **kwargs):
raise RuntimeError("Transonic types are not meant to be instantiated")
class MemLayout(Enum):
C = auto()
F = auto()
C_or_F = auto()
strided = auto()
def __repr__(self):
return f'"{self.name}"'
class ArrayMeta(Meta):
......@@ -216,8 +228,11 @@ class ArrayMeta(Meta):
dtype = None
ndim = None
memview = False
mem_layout = MemLayout.C_or_F
params_filtered = []
for param in parameters:
if param is None:
continue
if isinstance(param, (Type, type, np.dtype)):
if dtype is not None:
raise ValueError(
......@@ -263,6 +278,11 @@ class ArrayMeta(Meta):
if param == "memview":
memview = True
continue
try:
mem_layout = MemLayout[param]
continue
except KeyError:
pass
raise ValueError(f"{param} cannot be interpretted...")
params_filtered.append(param)
......@@ -281,6 +301,7 @@ class ArrayMeta(Meta):
{"dtype": dtype, "ndim": ndim, "parameters": parameters},
)
ArrayBis.memview = memview
ArrayBis.mem_layout = mem_layout
return ArrayBis
......@@ -309,6 +330,9 @@ class ArrayMeta(Meta):
if self.memview:
strings.append('"memview"')
if self.mem_layout is not MemLayout.C_or_F:
strings.append(repr(self.mem_layout))
return f"Array[{', '.join(strings)}]"
def format_as_backend_type(self, backend_type_formatter, **kwargs):
......@@ -340,7 +364,9 @@ class ArrayMeta(Meta):
raise ValueError
memview = kwargs.get("memview", self.memview)
return backend_type_formatter.make_array_code(dtype, ndim, memview)
return backend_type_formatter.make_array_code(
dtype, ndim, memview, self.mem_layout
)
class Array(metaclass=ArrayMeta):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment