Move some symbols to the root, fix docs

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",
+}

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",

loopy/codegen/__init__.py

@@ -87,6 +87,12 @@
 .. automodule:: loopy.codegen.result
 
 .. automodule:: loopy.codegen.tools
+
+References
+^^^^^^^^^^
+.. class:: Expression
+
+    See :class:`pymbolic.Expression`.
 """
 
 

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

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,40 +83,37 @@
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 .. 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
 
 
 Expression Manipulation Helpers
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 .. 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):

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,