Add docs for FucntionIdT, CallablesTable

illinois-ceesd · Apr 24, 2024 · 3666639 · 3666639
1 parent e999af6
commit 3666639
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 15 deletions.
diff --git a/doc/ref_call.rst b/doc/ref_call.rst
@@ -26,8 +26,8 @@ call site. For example, a call to ``sin(x)`` in :mod:`loopy` is type-generic to
 begin with, but it later specialized to either ``sinf``, ``sin`` or ``sinl``
 depending on the type of its argument ``x``. A callable's behavior during type
 or shape specialization is encoded via
-:meth:`~loopy.kernel.function_interface.InKernelCallable.with_types` and
-:meth:`~loopy.kernel.function_interface.InKernelCallable.with_descrs`.
+:meth:`~loopy.InKernelCallable.with_types` and
+:meth:`~loopy.InKernelCallable.with_descrs`.
 
 
 Registering callables

diff --git a/doc/ref_translation_unit.rst b/doc/ref_translation_unit.rst
@@ -1,11 +1,6 @@
 .. currentmodule:: loopy
 
-TranslationUnit
-===============
-
-.. autoclass:: TranslationUnit
-
-Reference
----------
+Translation Units
+=================
 
 .. automodule:: loopy.translation_unit
diff --git a/loopy/__init__.py b/loopy/__init__.py
@@ -45,7 +45,7 @@
         SubstitutionRule,
         CallMangleInfo)
 from loopy.kernel.function_interface import (
-        CallableKernel, ScalarCallable)
+        InKernelCallable, CallableKernel, ScalarCallable)
 from loopy.translation_unit import (
         TranslationUnit, make_program)
 
@@ -186,7 +186,7 @@
         "CallInstruction", "CInstruction", "NoOpInstruction",
         "BarrierInstruction",
 
-        "ScalarCallable", "CallableKernel",
+        "InKernelCallable", "ScalarCallable", "CallableKernel",
 
         "TranslationUnit", "make_program",
 

diff --git a/loopy/kernel/function_interface.py b/loopy/kernel/function_interface.py
@@ -43,6 +43,8 @@
 
 .. autoclass:: ArrayArgDescriptor
 
+.. currentmodule:: loopy
+
 .. autoclass:: InKernelCallable
 
 .. autoclass:: CallableKernel
@@ -69,7 +71,7 @@ class ArrayArgDescriptor(ImmutableRecord):
     """
     Records information about an array argument to an in-kernel callable. To be
     passed to and returned from
-    :meth:`InKernelCallable.with_descrs`, used for
+    :meth:`~loopy.InKernelCallable.with_descrs`, used for
     matching shape and address space of caller and callee kernels.
 
     .. attribute:: shape
@@ -372,9 +374,10 @@ def with_types(self, arg_id_to_dtype, clbl_inf_ctx):
     def with_descrs(self, arg_id_to_descr, clbl_inf_ctx):
         """
         :arg arg_id_to_descr: a mapping from argument identifiers (integers for
-            positional arguments) to instances of :class:`ArrayArgDescriptor`
-            or :class:`ValueArgDescriptor`. Unspecified/unknown descriptors are
-            not represented in *arg_id_to_type*.
+            positional arguments) to instances of
+            :class:`~loopy.kernel.function_interface.ArrayArgDescriptor`
+            or :class:`~loopy.kernel.function_interface.ValueArgDescriptor`.
+            Unspecified/unknown descriptors are not represented in *arg_id_to_type*.
 
             Return values are denoted by negative integers, with the first
             returned value identified as *-1*.

diff --git a/loopy/translation_unit.py b/loopy/translation_unit.py
@@ -48,6 +48,21 @@
 
 
 __doc__ = """
+
+.. class:: FunctionIdT
+
+    A type for a function identifier.
+    A :class:`~loopy.library.reduction.ReductionOpFunction` or a :class:`str`.
+
+.. class:: CallablesTable
+
+    A type alias for callables tables, mapping from :class:`FunctionIdT`
+    to :class:`~loopy.InKernelCallable`
+
+.. currentmodule:: loopy
+
+.. autoclass:: TranslationUnit
+
 .. currentmodule:: loopy.translation_unit
 
 .. autoclass:: CallablesInferenceContext