From fce1938f5ef9835e5967b17540af27a0c6bb030f Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Thu, 3 Oct 2024 10:02:12 -0500 Subject: [PATCH] Move some symbols to the root, fix docs --- doc/conf.py | 8 ++++++++ loopy/__init__.py | 5 ++++- loopy/codegen/__init__.py | 6 ++++++ loopy/kernel/data.py | 7 +++++++ loopy/symbolic.py | 42 +++++++++++++++++++-------------------- loopy/types.py | 9 +++++++-- 6 files changed, 52 insertions(+), 25 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index 2369d8b59..d14336d29 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -61,3 +61,11 @@ # available. ["py:class", r"immutables\.(.+)"], ] + +autodoc_type_aliases = { + "ToLoopyTypeConvertible": "ToLoopyTypeConvertible", + "ExpressionT": "ExpressionT", + "InameStr": "InameStr", + "ShapeType": "ShapeType", + "StridesType": "StridesType", +} diff --git a/loopy/__init__.py b/loopy/__init__.py index 149a1af26..07f06a021 100644 --- a/loopy/__init__.py +++ b/loopy/__init__.py @@ -205,7 +205,7 @@ ) from loopy.translation_unit import TranslationUnit, for_each_kernel, make_program from loopy.type_inference import infer_unknown_types -from loopy.types import to_loopy_type +from loopy.types import LoopyType, NumpyType, ToLoopyTypeConvertible, to_loopy_type from loopy.typing import auto from loopy.version import MOST_RECENT_LANGUAGE_VERSION, VERSION @@ -248,12 +248,14 @@ "LinearSubscript", "LoopKernel", "LoopyError", + "LoopyType", "LoopyWarning", "MemAccess", "MemoryOrdering", "MemoryScope", "MultiAssignmentBase", "NoOpInstruction", + "NumpyType", "Op", "OpenCLTarget", "Optional", @@ -270,6 +272,7 @@ "TemporaryVariable", "ToCountMap", "ToCountPolynomialMap", + "ToLoopyTypeConvertible", "TranslationUnit", "TypeCast", "UniqueName", diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index 268e54ec1..2e39d89bd 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -87,6 +87,12 @@ .. automodule:: loopy.codegen.result .. automodule:: loopy.codegen.tools + +References +^^^^^^^^^^ +.. class:: Expression + + See :class:`pymbolic.Expression`. """ diff --git a/loopy/kernel/data.py b/loopy/kernel/data.py index b6cc601c0..bff69a977 100644 --- a/loopy/kernel/data.py +++ b/loopy/kernel/data.py @@ -87,6 +87,13 @@ .. autoclass:: UnrollTag .. autoclass:: Iname + +References +^^^^^^^^^^ + +.. class:: ToLoopyTypeConvertible + + See :class:`loopy.ToLoopyTypeConvertible`. """ # This docstring is included in ref_internals. Do not include parts of the public diff --git a/loopy/symbolic.py b/loopy/symbolic.py index 71a0c777e..989a741da 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -40,6 +40,7 @@ import pymbolic.primitives as p import pytools.lex from islpy import dim_type +from pymbolic import Variable from pymbolic.mapper import ( CachedCombineMapper as CombineMapperBase, CachedIdentityMapper as IdentityMapperBase, @@ -82,33 +83,23 @@ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. autoclass:: Literal - .. autoclass:: ArrayLiteral - .. autoclass:: FunctionIdentifier - .. autoclass:: TypedCSE .. currentmodule:: loopy .. autoclass:: TypeCast - .. autoclass:: TaggedVariable - .. autoclass:: Reduction - .. autoclass:: LinearSubscript .. currentmodule:: loopy.symbolic .. autoclass:: RuleArgument - .. autoclass:: ExpansionState - .. autoclass:: RuleAwareIdentityMapper - .. autoclass:: ResolvedFunction - .. autoclass:: SubArrayRef @@ -116,6 +107,13 @@ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. autofunction:: simplify_using_aff + +References +^^^^^^^^^^ + +.. class:: Variable + + See :class:`pymbolic.Variable`. """ @@ -144,7 +142,7 @@ def map_reduction(self, expr, *args, **kwargs): new_inames = [] for iname, new_sym_iname in zip(expr.inames, mapped_inames): - if not isinstance(new_sym_iname, p.Variable): + if not isinstance(new_sym_iname, Variable): from loopy.diagnostic import LoopyError raise LoopyError("%s did not map iname '%s' to a variable" % (type(self).__name__, iname)) @@ -412,7 +410,7 @@ def map_call(self, expr, *args, **kwargs): def map_reduction(self, expr, *args, **kwargs): deps = self.rec(expr.expr, *args, **kwargs) - return deps - {p.Variable(iname) for iname in expr.inames} + return deps - {Variable(iname) for iname in expr.inames} def map_tagged_variable(self, expr, *args, **kwargs): return {expr} @@ -625,7 +623,7 @@ def type(self) -> NumpyType: @p.expr_dataclass(init=False) -class TaggedVariable(LoopyExpressionBase, p.Variable, Taggable): +class TaggedVariable(LoopyExpressionBase, Variable, Taggable): """This is an identifier with tags, such as ``matrix$one``, where 'one' identifies this specific use of the identifier. This mechanism may then be used to address these uses--such as by prefetching only @@ -646,7 +644,7 @@ class TaggedVariable(LoopyExpressionBase, p.Variable, Taggable): """ def __init__(self, name: str, tags: ToTagSetConvertible) -> None: - p.Variable.__init__(self, name) + Variable.__init__(self, name) if isinstance(tags, str): from loopy.kernel.creation import _normalize_string_tag tags = frozenset({_normalize_string_tag(tags)}) @@ -706,13 +704,13 @@ def __init__(self, if isinstance(inames, str): inames = tuple(iname.strip() for iname in inames.split(",")) - elif isinstance(inames, p.Variable): + elif isinstance(inames, Variable): inames = (inames,) assert isinstance(inames, tuple) def strip_var(iname: Any) -> str: - if isinstance(iname, p.Variable): + if isinstance(iname, Variable): iname = iname.name assert isinstance(iname, str) @@ -787,19 +785,19 @@ class ResolvedFunction(LoopyExpressionBase): .. autoattribute:: function .. autoattribute:: name """ - function: p.Variable | ReductionOpFunction + function: Variable | ReductionOpFunction def __init__(self, function): if isinstance(function, str): - function = p.Variable(function) + function = Variable(function) from loopy.library.reduction import ReductionOpFunction - assert isinstance(function, (p.Variable, ReductionOpFunction)) + assert isinstance(function, (Variable, ReductionOpFunction)) object.__setattr__(self, "function", function) @property def name(self): from loopy.library.reduction import ReductionOpFunction - if isinstance(self.function, p.Variable): + if isinstance(self.function, Variable): return self.function.name elif isinstance(self.function, ReductionOpFunction): return self.function @@ -827,7 +825,7 @@ def map_variable(self, expr): class VariableInAnExpression(CombineMapper): def __init__(self, variables_to_search): - assert all(isinstance(variable, p.Variable) for variable in + assert all(isinstance(variable, Variable) for variable in variables_to_search) self.variables_to_search = variables_to_search @@ -895,7 +893,7 @@ class SubArrayRef(LoopyExpressionBase): .. automethod:: is_equal """ - swept_inames: tuple[p.Variable, ...] + swept_inames: tuple[Variable, ...] subscript: p.Subscript def __post_init__(self): diff --git a/loopy/types.py b/loopy/types.py index a837d1c46..b43026bdb 100644 --- a/loopy/types.py +++ b/loopy/types.py @@ -26,18 +26,23 @@ from typing import Any, Mapping, Type, Union import numpy as np +from typing_extensions import TypeAlias from loopy.diagnostic import LoopyError from loopy.typing import auto __doc__ = """ -.. currentmodule:: loopy.types +.. currentmodule:: loopy .. autoclass:: LoopyType +.. autoclass:: ToLoopyTypeConvertible + .. autoclass:: NumpyType +.. currentmodule:: loopy.types + .. autoclass:: AtomicType .. autoclass:: AtomicNumpyType @@ -197,7 +202,7 @@ def __eq__(self, other: object) -> bool: # }}} -ToLoopyTypeConvertible = Union[Type[auto], None, np.dtype, LoopyType] +ToLoopyTypeConvertible: TypeAlias = Union[Type[auto], None, np.dtype, LoopyType] def to_loopy_type(dtype: ToLoopyTypeConvertible,