From 57f15e6ac9a389d140aecb2e8bf7656fbbf43bad Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Thu, 19 Dec 2024 11:48:59 -0600 Subject: [PATCH] Fix up the doc build --- doc/conf.py | 2 ++ doc/ref_kernel.rst | 8 +------- loopy/codegen/__init__.py | 18 ++++++++++++----- loopy/kernel/__init__.py | 41 +++++++++++++++++++++++++-------------- loopy/kernel/array.py | 2 +- loopy/typing.py | 4 ++++ 6 files changed, 47 insertions(+), 28 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index 70f7121b4..59d9a2c94 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -63,6 +63,8 @@ # Reference not found from ""? I'm not even sure where to look. ["py:class", r"ExpressionNode"], + + ["py:class", r"InameStr"], ] autodoc_type_aliases = { diff --git a/doc/ref_kernel.rst b/doc/ref_kernel.rst index 02f8b1485..227356b11 100644 --- a/doc/ref_kernel.rst +++ b/doc/ref_kernel.rst @@ -656,8 +656,6 @@ Helper values .. {{{ -.. autoclass:: auto - .. autoclass:: UniqueName .. autoclass:: Optional @@ -693,11 +691,7 @@ The Kernel Object Do not create :class:`LoopKernel` objects directly. Instead, refer to :ref:`creating-kernels`. -.. autoclass:: LoopKernel - -.. autoclass:: KernelState - :members: - :undoc-members: +.. automodule:: loopy.kernel Implementation Details: The Base Array -------------------------------------- diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index bc4ec29bd..2e053fd4f 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -33,7 +33,7 @@ Sequence, ) -from immutables import Map +import immutables from loopy.codegen.result import CodeGenerationResult from loopy.library.reduction import ReductionOpFunction @@ -44,8 +44,10 @@ from functools import reduce +import islpy # to help out Sphinx import islpy as isl -from pytools import ProcessLogger, UniqueNameGenerator +import pytools # to help out Sphinx +from pytools import ProcessLogger from pytools.persistent_dict import WriteOncePersistentDict from loopy.diagnostic import LoopyError, warn @@ -157,7 +159,7 @@ class CodeGenerationState: """ kernel: LoopKernel - implemented_domain: isl.Set + implemented_domain: islpy.Set """ The entire implemented domain (as an :class:`islpy.Set`) i.e. all constraints that have been enforced so far. @@ -169,11 +171,11 @@ class CodeGenerationState: seen_functions: set[SeenFunction] seen_atomic_dtypes: set[LoopyType] - var_subst_map: Map[str, Expression] + var_subst_map: immutables.Map[str, Expression] allow_complex: bool callables_table: CallablesTable is_entrypoint: bool - var_name_generator: UniqueNameGenerator + var_name_generator: pytools.UniqueNameGenerator is_generating_device_code: bool gen_program_name: str @@ -328,6 +330,12 @@ def map_constant(self, expr): @dataclass(frozen=True) class PreambleInfo: + """ + .. autoattribute:: kernel + .. autoattribute:: seen_dtypes + .. autoattribute:: seen_functions + .. autoattribute:: seen_atomic_dtypes + """ kernel: LoopKernel seen_dtypes: set[LoopyType] seen_functions: set[SeenFunction] diff --git a/loopy/kernel/__init__.py b/loopy/kernel/__init__.py index b3beb9336..ca4e6452a 100644 --- a/loopy/kernel/__init__.py +++ b/loopy/kernel/__init__.py @@ -1,4 +1,12 @@ -"""Kernel object.""" +""" +.. currentmodule:: loopy + +.. autoclass:: LoopKernel + +.. autoclass:: KernelState + :members: + :undoc-members: +""" from __future__ import annotations @@ -23,7 +31,6 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ - from collections import defaultdict from dataclasses import dataclass, field, fields, replace from enum import IntEnum @@ -43,7 +50,9 @@ import numpy as np from immutables import Map +from typing_extensions import TypeAlias +import islpy # to help out Sphinx import islpy as isl from islpy import dim_type from pymbolic import ArithmeticExpression @@ -55,10 +64,10 @@ ) from pytools.tag import Tag, Taggable +import loopy.kernel.data # to help out Sphinx from loopy.diagnostic import CannotBranchDomainTree, LoopyError, StaticValueFindingError from loopy.kernel.data import ( ArrayArg, - Iname, KernelArgument, SubstitutionRule, TemporaryVariable, @@ -76,7 +85,7 @@ if TYPE_CHECKING: - from loopy.codegen import PreambleInfo + import loopy.codegen # to help out Sphinx from loopy.kernel.function_interface import InKernelCallable @@ -96,12 +105,9 @@ def _get_inames_from_domains(domains): @dataclass(frozen=True) class _BoundsRecord: - lower_bound_pw_aff: isl.PwAff - upper_bound_pw_aff: isl.PwAff - size: isl.PwAff - - -PreambleGenerator = Callable[["PreambleInfo"], Iterator[Tuple[int, str]]] + lower_bound_pw_aff: islpy.PwAff + upper_bound_pw_aff: islpy.PwAff + size: islpy.PwAff @dataclass(frozen=True) @@ -141,7 +147,7 @@ class LoopKernel(Taggable): .. automethod:: tagged .. automethod:: without_tags """ - domains: Sequence[isl.BasicSet] + domains: Sequence[islpy.BasicSet] """Represents the :ref:`domain-tree`.""" instructions: Sequence[InstructionBase] @@ -150,13 +156,13 @@ class LoopKernel(Taggable): """ args: Sequence[KernelArgument] - assumptions: isl.BasicSet + assumptions: islpy.BasicSet """ Must be a :class:`islpy.BasicSet` parameter domain. """ temporary_variables: Mapping[str, TemporaryVariable] - inames: Mapping[InameStr, Iname] + inames: Mapping[InameStr, loopy.kernel.data.Iname] """ An entry is guaranteed to be present for each iname. """ @@ -169,7 +175,11 @@ class LoopKernel(Taggable): name: str = "loopy_kernel" preambles: Sequence[tuple[int, str]] = () - preamble_generators: Sequence[PreambleGenerator] = () + preamble_generators: Sequence[ + Callable[ + ["loopy.codegen.PreambleInfo"], + Iterator[Tuple[int, str]]] + ] = () symbol_manglers: Sequence[ Callable[[LoopKernel, str], tuple[LoopyType, str] | None]] = () linearization: Sequence[ScheduleItem] | None = None @@ -1360,7 +1370,8 @@ def get_copy_kwargs(self, **kwargs: Any) -> dict[str, Any]: if "domains" in kwargs: inames = kwargs.get("inames", self.inames) domains = kwargs["domains"] - kwargs["inames"] = {name: inames.get(name, Iname(name, frozenset())) + kwargs["inames"] = {name: inames.get(name, + loopy.kernel.data.Iname(name, frozenset())) for name in _get_inames_from_domains(domains)} assert all(dom.get_ctx() == isl.DEFAULT_CONTEXT for dom in domains) diff --git a/loopy/kernel/array.py b/loopy/kernel/array.py index 240c74e13..a91f7925b 100644 --- a/loopy/kernel/array.py +++ b/loopy/kernel/array.py @@ -55,7 +55,7 @@ if TYPE_CHECKING: from loopy.codegen import VectorizationInfo from loopy.kernel import LoopKernel - from loopy.kernel.data import ArrayArg, TemporaryVariable, auto + from loopy.kernel.data import ArrayArg, TemporaryVariable from loopy.target import TargetBase if getattr(sys, "_BUILDING_SPHINX_DOCS", False): diff --git a/loopy/typing.py b/loopy/typing.py index 3260c0f41..9b11720f4 100644 --- a/loopy/typing.py +++ b/loopy/typing.py @@ -1,6 +1,10 @@ """ .. autoclass:: Expression .. autoclass:: ShapeType +.. autodata:: InameStr + +.. currentmodule:: loopy + .. autoclass:: auto """