This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new cbca9548f [Rust] Native docs for tirx, sym, backend and the docs pages
(#221)
cbca9548f is described below
commit cbca9548fb8ec1ef4a7219dcaa6a5ecaa0b1535c
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Nov 6 22:31:22 2026 -0500
[Rust] Native docs for tirx, sym, backend and the docs pages (#221)
Continue the native rustdoc pass from #220 over the rest of the crate
and the Rust docs pages. Docs only; no code change.
- `tirx`, `sym`, `backend`, `runtime`, `target`, `codegen`, and `script`
modules get short docs on what each is about and how its central types
relate; examples move from module docs to the key types and functions.
- The object types (43) get type docs on the reference types, written
above the generated blocks so regeneration keeps them.
- Rustdoc no longer mentions C++; method docs describe behavior in Rust
terms.
- The crate doc gets a rewritten module list and drops the claims that
constructors follow the C++ ones and modules follow the Python package.
- `defining_ir_extensions.rst` and `working_with_ir_in_rust.rst` drop
C++ comparisons, and the `#[dispatch]` quick reference matches the
current signature form.
---
docs/advanced/defining_ir_extensions.rst | 33 +++++-----
docs/ir_guides/working_with_ir_in_rust.rst | 10 ++-
rust/src/backend.rs | 3 +-
rust/src/backend/cuda.rs | 2 +-
rust/src/backend/cuda/op.rs | 18 ++++--
rust/src/codegen.rs | 5 +-
rust/src/codegen/llvm.rs | 39 ++++++------
rust/src/lib.rs | 30 +++++----
rust/src/macro/def_expr_op_overloads.rs | 6 +-
rust/src/runtime.rs | 41 +++++++------
rust/src/script.rs | 45 +++++++-------
rust/src/sym.rs | 7 +--
rust/src/sym/analyzer.rs | 71 ++++++++++++++--------
rust/src/sym/constraint.rs | 63 ++++++++++---------
rust/src/sym/int_set.rs | 49 ++++++++-------
rust/src/sym/iter_affine_map.rs | 68 ++++++++++++---------
rust/src/target.rs | 44 ++++++++------
rust/src/tirx.rs | 7 +--
rust/src/tirx/analysis.rs | 4 +-
rust/src/tirx/analysis/analysis.rs | 32 +++++-----
rust/src/tirx/function.rs | 54 ++++++++--------
rust/src/tirx/index_map.rs | 53 ++++++++--------
rust/src/tirx/layout.rs | 67 ++++++++++++--------
rust/src/tirx/op.rs | 52 ++++++++--------
rust/src/tirx/op_attr.rs | 8 +--
rust/src/tirx/structural.rs | 15 +++--
.../tirx/structural_map_with_constraint_context.rs | 4 +-
.../structural_mutate_with_constraint_context.rs | 4 +-
.../structural_visit_with_constraint_context.rs | 4 +-
.../structural_walk_with_constraint_context.rs | 3 +-
rust/src/tirx/tensor.rs | 69 +++++++++++----------
rust/src/tirx/thread_scope.rs | 35 ++++++-----
rust/src/tirx/tile.rs | 4 +-
rust/src/tirx/tile/dispatch.rs | 63 ++++++++++---------
rust/src/tirx/tile/instruction.rs | 11 ++--
rust/src/tirx/transform.rs | 3 +-
rust/src/tirx/transform/function_pass.rs | 57 ++++++++---------
rust/src/tirx/transform/transform.rs | 57 +++++++++--------
rust/src/tirx/type.rs | 22 ++-----
39 files changed, 628 insertions(+), 534 deletions(-)
diff --git a/docs/advanced/defining_ir_extensions.rst
b/docs/advanced/defining_ir_extensions.rst
index f69099318..9a872488c 100644
--- a/docs/advanced/defining_ir_extensions.rst
+++ b/docs/advanced/defining_ir_extensions.rst
@@ -77,9 +77,8 @@ real and imaginary parts, as in ``TensorType``:
``#[derive(ObjectRef)]`` makes ``ComplexType``, whose one field is an
``ObjectArc<ComplexTypeObj>``, the reference type of the object type
- ``ComplexTypeObj``, as the C++ macro
- ``TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE`` does. A reference holds a
- reference-counted pointer to its object, and the derive gives it:
+ ``ComplexTypeObj``. A reference holds a reference-counted pointer to its
+ object, and the derive gives it:
- ``Deref`` to the object, so its fields read directly: ``ty.dtype``;
- a private ``from_object``, which wraps a new ``ComplexTypeObj`` for the
@@ -89,12 +88,12 @@ real and imaginary parts, as in ``TensorType``:
``into()`` wherever one is expected, such as the result of type
inference.
-In Rust, as in C++, the type is defined in two parts: ``ComplexTypeObj`` is the
-object that holds the fields, and ``ComplexType`` is a reference-counted
-reference to it, which IR nodes store and functions pass around.
-``base_: TypeObj`` makes the object an IR type, and ``rs_object`` registers its
-fields for reflection, through which Python and IR JSON loading construct the
-type as they construct a C++ one. ``#[derive(ObjectRef)]`` defines the
reference.
+In Rust, the type is defined in two parts: ``ComplexTypeObj`` is the object
that
+holds the fields, and ``ComplexType`` is a reference-counted reference to it,
+which IR nodes store and functions pass around. ``base_: TypeObj`` makes the
+object an IR type, and ``rs_object`` registers its fields for reflection,
+through which Python and IR JSON loading construct the type as they construct
+a built-in one. ``#[derive(ObjectRef)]`` defines the reference.
The ``static_init_block!`` registers the type when the library loads.
@@ -103,21 +102,21 @@ The ``static_init_block!`` registers the type when the
library loads.
**What does static_init_block! do?**
``static_init_block! { ... }`` runs its body when the program or library
that
- holds it loads, as a C++ ``TVM_FFI_STATIC_INIT_BLOCK`` does. The body
- registers what the library defines, such as types, operators, and global
- functions, and Python sees these registrations once it loads the library.
+ holds it loads. The body registers what the library defines, such as types,
+ operators, and global functions, and Python sees these registrations once it
+ loads the library.
The body is the body of a function that returns ``Result<()>``, so ``?``
stops it at its first error.
Here, ``register_type::<ComplexType>()`` registers ``ComplexType`` as an IR
type, so that Python constructs it from its fields and IR JSON holds it, as
- for a C++ type.
+ for a built-in type.
In Python, one class is enough. ``py_class`` from ``tvm_ffi.dataclasses``
registers it as an IR object: ``my_ext.ComplexType`` is its type key,
``ir.Type`` its parent, and the annotated ``dtype`` its field. The class gets a
constructor, ``ComplexType(dtype)``, and its field takes part in structural
-equality, printing, and IR JSON, as the fields of a C++ type do. The global
+equality, printing, and IR JSON, as the fields of a built-in type do. The
global
function ``ir.expose_type`` then publishes the type, as ``register_type`` does
in Rust.
@@ -267,9 +266,9 @@ that pair. This lets the rules apply to the step-by-step
function from
Im[i] = product_im
Only scalar arithmetic remains. To apply the rules, the pass tests whether a
-call's ``op`` is one of the four operators. The Rust pass looks them up by
name,
-as C++ does with ``Op::Get``. The operators are registered when the library
-loads, and ``Op::get`` returns the registry's single object for a name.
+call's ``op`` is one of the four operators. The Rust pass looks them up by
name.
+The operators are registered when the library loads, and ``Op::get`` returns
+the registry's single object for a name.
``call.op.same_as(&Op::get("my_ext.cmul")?)`` therefore tests whether a call's
``op`` is ``my_ext.cmul`` by comparing pointers. The Python pass compares with
the operators that ``Op.register`` returned, as in ``call.op.same_as(CMUL)``.
diff --git a/docs/ir_guides/working_with_ir_in_rust.rst
b/docs/ir_guides/working_with_ir_in_rust.rst
index 39f48ee37..32dc4ed3c 100644
--- a/docs/ir_guides/working_with_ir_in_rust.rst
+++ b/docs/ir_guides/working_with_ir_in_rust.rst
@@ -143,7 +143,7 @@ The output is:
directly from ``x``, and the literals ``2`` and ``1`` become constants of
its type.
-As in C++, the operators take a ``PrimExpr``, an expression whose type is a
+The operators take a ``PrimExpr``, an expression whose type is a
``PrimType``. ``x`` is a ``PrimVar``, a variable of a primitive type, because a
plain ``Var`` may have any type: it takes part in arithmetic only after
``downcast_ref::<PrimVar>()`` or ``downcast::<PrimVar>()`` checks it, and
``&var + 1``
@@ -215,8 +215,8 @@ This is the newtype pattern. Simplified from the crate, the
pair is:
Rust has no inheritance; instead, ``#[repr(C)]`` puts the parent's object,
``ExprObj``, first in ``AddObj``, so every addition is also an expression.
``Add`` implements ``Deref`` to ``AddObj``, which is why ``add.a`` works.
-``AddObj`` has the same memory layout as its C++ counterpart, so
-``add.a`` is a direct field read, as efficient as in C++.
+``AddObj`` has the same memory layout as the native object, so
+``add.a`` is a direct field read.
**Nodes are immutable.** As with data behind an ``Arc``, a managed
reference gives read access only; to change a node, we build a new one.
@@ -791,7 +791,7 @@ Quick Reference
- Receives the definition region kind
* - |structural_equal|_\ ``(&a, &b)``
- Compares by structure
- * - |dispatch|_\ ``(prefix, (Carrier, ...), ResultType)``
+ * - |dispatch|_\ ``(prefix, (First, Args...), Ret)``
- Generates ``dispatch``, a ``match_any!`` over the ``prefix_*``
methods of a struct
* - |structural_visit|_, |structural_mutate|_
@@ -848,8 +848,6 @@ Quick Reference
.. _structural_walk_with_constraint_context:
../reference/rust/generated/tvm_tirx/tirx/structural_walk_with_constraint_context/fn.structural_walk_with_constraint_context.html
.. |structural_mutate_with_constraint_context| replace::
``structural_mutate_with_constraint_context``
.. _structural_mutate_with_constraint_context:
../reference/rust/generated/tvm_tirx/tirx/structural_mutate_with_constraint_context/fn.structural_mutate_with_constraint_context.html
-.. |StructuralVisitorWithConstraintContextView| replace::
``StructuralVisitorWithConstraintContextView``
-.. _StructuralVisitorWithConstraintContextView:
../reference/rust/generated/tvm_tirx/tirx/structural_visit_with_constraint_context/type.StructuralVisitorWithConstraintContextView.html
.. |match_any!| replace:: ``match_any!``
.. _match_any!: ../reference/rust/generated/tvm_ffi/macro.match_any.html
.. |PrimVar::new| replace:: ``PrimVar::new``
diff --git a/rust/src/backend.rs b/rust/src/backend.rs
index 05299c460..0a2159aea 100644
--- a/rust/src/backend.rs
+++ b/rust/src/backend.rs
@@ -17,6 +17,7 @@
* under the License.
*/
-//! Backend-specific objects, such as the attributes of CUDA instructions in
[`cuda`].
+//! Backends. Each submodule holds the objects of one backend, such as the
attributes of the CUDA
+//! instructions in [`cuda`].
pub mod cuda;
diff --git a/rust/src/backend/cuda.rs b/rust/src/backend/cuda.rs
index 48263277b..ee3ef18e0 100644
--- a/rust/src/backend/cuda.rs
+++ b/rust/src/backend/cuda.rs
@@ -17,7 +17,7 @@
* under the License.
*/
-//! The CUDA backend: the objects of the `tirx.cuda` namespace.
+//! The CUDA backend. [`op`] holds the static attributes its instructions
carry.
// tvm-ffi-stubgen(missing-docs): allow
// An object registered under `tirx.cuda` that no file has a block for yet is
generated here, to be
diff --git a/rust/src/backend/cuda/op.rs b/rust/src/backend/cuda/op.rs
index 92d010a37..7c329abb8 100644
--- a/rust/src/backend/cuda/op.rs
+++ b/rust/src/backend/cuda/op.rs
@@ -21,10 +21,9 @@
// The generated constructors take one argument per reflected field.
#![allow(clippy::too_many_arguments)]
-//! The attributes of CUDA instructions.
-//!
-//! [`TCGen05InstrDescriptorAttrs`] and
[`TCGen05InstrDescriptorBlockScaledAttrs`] describe a
-//! `tcgen05` instruction descriptor, and [`TensorMapEncodeTiledAttr`] a tiled
tensor-map encoding.
+//! Attributes of CUDA instructions. [`TCGen05InstrDescriptorAttrs`] and
+//! [`TCGen05InstrDescriptorBlockScaledAttrs`] are the static options of a
`tcgen05` instruction
+//! descriptor, and [`TensorMapEncodeTiledAttr`] those of a tiled tensor-map
encoding.
// tvm-ffi-stubgen(missing-docs): allow
@@ -37,6 +36,10 @@ use tvm_ffi::ObjectArc;
use tvm_ffi::String;
// tvm-ffi-stubgen(end)
+/// The static options of a dense `tcgen05` instruction descriptor, which
+/// `tirx.cuda.tcgen05_encode_instr_descriptor` encodes: the data types of the
result `d_dtype` and
+/// the operands `a_dtype` and `b_dtype`, the shape `M` x `N` x `K`, the
operand transposes, the
+/// number of CTA groups, operand negation, result saturation `sat_d` and
sparsity.
// tvm-ffi-stubgen(begin): object/tirx.cuda.TCGen05InstrDescriptorAttrs
#[allow(missing_docs)]
#[repr(C)]
@@ -164,6 +167,10 @@ unsafe impl tvm_ffi::object::ImmutableObject for
TCGen05InstrDescriptorAttrsObj
tvm_ffi::unsafe_impl_object_ref_upcast!(TCGen05InstrDescriptorAttrs => Attrs);
// tvm-ffi-stubgen(end)
+/// The static options of a block-scaled `tcgen05` instruction descriptor,
which
+/// `tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled` encodes: those of
+/// [`TCGen05InstrDescriptorAttrs`] but saturation, and the data types of the
scale factors
+/// `sfa_dtype` and `sfb_dtype`.
// tvm-ffi-stubgen(begin):
object/tirx.cuda.TCGen05InstrDescriptorBlockScaledAttrs
#[allow(missing_docs)]
#[repr(C)]
@@ -296,6 +303,9 @@ unsafe impl tvm_ffi::object::ImmutableObject for
TCGen05InstrDescriptorBlockScal
tvm_ffi::unsafe_impl_object_ref_upcast!(TCGen05InstrDescriptorBlockScaledAttrs
=> Attrs);
// tvm-ffi-stubgen(end)
+/// The fixed options of a `tirx.cuda.tensormap_encode_tiled` call, which
encodes a tiled tensor map
+/// at run time: the `descriptor_dtype` and `rank` of the tensor, and its
interleave, swizzle, L2
+/// promotion, out-of-bounds fill and data-type options.
// tvm-ffi-stubgen(begin): object/tirx.cuda.TensorMapEncodeTiledAttr
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/codegen.rs b/rust/src/codegen.rs
index 102166cc3..ac57ae3f0 100644
--- a/rust/src/codegen.rs
+++ b/rust/src/codegen.rs
@@ -17,9 +17,8 @@
* under the License.
*/
-//! Code generation: IR modules compiled into runtime modules.
-//!
-//! The queries of the LLVM backend are in [`llvm`].
+//! Code generation. [`build_module`] compiles an [`IRModule`] into a runtime
[`Module`] with the
+//! code generator of the target's kind, and [`llvm`] queries the LLVM backend.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
diff --git a/rust/src/codegen/llvm.rs b/rust/src/codegen/llvm.rs
index 226552290..6b41fc5a9 100644
--- a/rust/src/codegen/llvm.rs
+++ b/rust/src/codegen/llvm.rs
@@ -17,26 +17,9 @@
* under the License.
*/
-//! The LLVM backend: intrinsics, the host CPU and the CPUs and features of
LLVM targets.
-//!
-//! The queries are registered by `libtvm_tirx_llvm_backend` when it is
loaded, as `import tvm_tirx`
-//! loads it. A query given no target uses the current one.
-//!
-//! # Examples
-//!
-//! ```
-//! use std::env::consts::{DLL_PREFIX, DLL_SUFFIX};
-//! use tvm_tirx::codegen::llvm;
-//! use tvm_tirx::target::Target;
-//! use tvm_tirx::tvm_ffi::Module;
-//!
-//! let backend = format!("{DLL_PREFIX}tvm_tirx_llvm_backend{DLL_SUFFIX}");
-//! let _backend = Module::load_from_file(backend)?;
-//! let cpu = llvm::get_system_cpu()?;
-//! let host = Target::new(&format!(r#"{{"kind": "llvm", "mcpu": "{}"}}"#,
cpu.as_str()))?;
-//! assert!(llvm::get_vector_width(Some(&host))? > 0);
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! The LLVM backend. Queries of LLVM intrinsics, the host CPU and the CPUs
and features of LLVM
+//! targets, which `libtvm_tirx_llvm_backend` registers when it is loaded; a
query given no target
+//! uses the current one.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -124,6 +107,22 @@ pub fn cpu_has_feature(feature: &str, target:
Option<&Target>) -> Result<bool> {
/// Return the vector width of the CPU of `target` in bits, or -1 for a target
that is not `llvm`
/// (`codegen.llvm_get_vector_width`).
+///
+/// # Examples
+///
+/// ```
+/// use std::env::consts::{DLL_PREFIX, DLL_SUFFIX};
+/// use tvm_tirx::codegen::llvm;
+/// use tvm_tirx::target::Target;
+/// use tvm_tirx::tvm_ffi::Module;
+///
+/// let backend = format!("{DLL_PREFIX}tvm_tirx_llvm_backend{DLL_SUFFIX}");
+/// let _backend = Module::load_from_file(backend)?;
+/// let cpu = llvm::get_system_cpu()?;
+/// let host = Target::new(&format!(r#"{{"kind": "llvm", "mcpu": "{}"}}"#,
cpu.as_str()))?;
+/// assert!(llvm::get_vector_width(Some(&host))? > 0);
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
pub fn get_vector_width(target: Option<&Target>) -> Result<i32> {
tvm_ffi::Function::get_global("codegen.llvm_get_vector_width")?
.call_tuple((target.cloned(),))?
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 27ac9cf21..4fc2cf726 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -20,17 +20,17 @@
//! The TIRx compiler in Rust: its IR objects over TVM FFI, and the definition
of new types,
//! operators and passes that extend it.
//!
-//! The IR bindings are generated by `tvm-ffi-stubgen` from the reflection
metadata of
-//! `libtvm_tirx_compiler`, and the modules follow the Python package
`tvm_tirx`:
+//! The modules:
//!
-//! - [`ir`]: the shared IR, with the primitive expressions in [`ir::prim`]
and the passes in
-//! [`ir::transform`].
-//! - [`tirx`]: tensor functions and types, layouts, traversals, tile
instructions and passes.
-//! - [`sym`]: symbolic analysis.
+//! - [`ir`]: the IR shared by every dialect, with the primitive expressions
in [`ir::prim`] and
+//! the pass infrastructure in [`ir::transform`].
+//! - [`tirx`]: tensor programs: functions, tensor types, layouts, tile
instructions, passes and
+//! analyses.
+//! - [`sym`]: symbolic integer analysis.
//! - [`target`]: compilation targets.
-//! - [`backend`]: backend-specific objects, such as the attributes of CUDA
instructions.
-//! - [`runtime`]: runtime objects, tensor allocation and runtime queries.
-//! - [`codegen`]: compilation of IR modules into runtime modules.
+//! - [`backend`]: backend objects, such as the attributes of CUDA
instructions.
+//! - [`runtime`]: runtime objects and queries, and tensor allocation.
+//! - [`codegen`]: compiling IR modules into runtime modules.
//! - [`script`]: printing IR objects as TVMScript.
//!
//! # Using the bindings
@@ -55,9 +55,8 @@
//! ```
//!
//! A handle such as [`ir::Expr`] is a reference-counted pointer to its
object, [`ir::ExprObj`], so
-//! cloning it does not copy the IR. The constructors (`new`) follow the C++
ones, with their checks
-//! and defaults, and return a `Result` when the C++ one can throw;
`from_fields` stores the fields
-//! as given.
+//! cloning it does not copy the IR. A constructor `new` checks its arguments
and fills in defaults,
+//! and returns a `Result` when a check can fail; `from_fields` stores the
fields as given.
//!
//! # Defining operators and types
//!
@@ -99,9 +98,8 @@
//! # Ok::<(), tvm_ffi::Error>(())
//! ```
//!
-//! A type extends [`ir::TypeObj`]; its fields are reflected, so C++, Python
and the TVMScript
-//! printer see them. The generated code names `tvm_ffi`, which `use
tvm_tirx::tvm_ffi` brings into
-//! scope.
+//! A type extends [`ir::TypeObj`]; its fields are reflected, so Python and
the TVMScript printer
+//! see them. The generated code names `tvm_ffi`, which `use
tvm_tirx::tvm_ffi` brings into scope.
//!
//! ```
//! use tvm_tirx::ir::{PrimType, Type, TypeObj, UNKNOWN_LOC};
@@ -138,7 +136,7 @@
//! # The library Python loads
//!
//! The crate is also built as `libtvm_tirx_rust`, which `import tvm_tirx`
loads, so Python uses the
-//! operators and types the crate defines as it uses the C++ ones. A `cdylib`
of operators of its
+//! operators and types the crate defines as it uses the native ones. A
`cdylib` of operators of its
//! own, such as `examples/define_ir.rs`, registers them when Python loads it;
a crate it depends on
//! but never names is not linked in, so name it, as `use my_ops as _;` does.
//!
diff --git a/rust/src/macro/def_expr_op_overloads.rs
b/rust/src/macro/def_expr_op_overloads.rs
index c8637709b..d3ea12b9c 100644
--- a/rust/src/macro/def_expr_op_overloads.rs
+++ b/rust/src/macro/def_expr_op_overloads.rs
@@ -23,9 +23,9 @@
/// The type, a module-scope path to a reference wrapper implementing `Clone`
and
/// `Into<PrimExpr>`, gets the operators `+`, `-`, `*`, `&`, `|`, `^`, `<<`,
`>>`, unary `-` and `!`
/// and [`PrimOps`](crate::ir::prim::PrimOps), all returning an
-/// [`ExprResult`](crate::ir::prim::ExprResult). As the C++ operators take a
`PrimExpr`, the type
-/// is one: an upcast to [`PrimExpr`](crate::ir::PrimExpr) declares that every
value of it is typed
-/// by a `PrimType`.
+/// [`ExprResult`](crate::ir::prim::ExprResult). The operators take
+/// [`PrimExpr`](crate::ir::PrimExpr) operands, so the type must convert into
one: its upcast to
+/// `PrimExpr` declares that every value of it is typed by a `PrimType`.
///
/// # Examples
///
diff --git a/rust/src/runtime.rs b/rust/src/runtime.rs
index 64990ffb1..68f3b78e5 100644
--- a/rust/src/runtime.rs
+++ b/rust/src/runtime.rs
@@ -17,25 +17,8 @@
* under the License.
*/
-//! The runtime objects the IR refers to, tensor allocation and runtime
queries.
-//!
-//! A [`MemoryScope`] is the storage scope of a
[`PtrType`](crate::ir::PtrType).
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::dtype;
-//! use tvm_tirx::runtime;
-//! use tvm_tirx::tvm_ffi::{DLDevice, DLDeviceType};
-//!
-//! let cpu = DLDevice::new(DLDeviceType::kDLCPU, 0);
-//! let buffer = runtime::empty(&[2, 8], dtype::float(32), cpu, None)?;
-//!
-//! let bytes = runtime::create_view(&buffer, &[64], dtype::int(8), 0)?;
-//! assert_eq!(bytes.shape(), &[64]);
-//!
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Runtime objects and queries. A [`MemoryScope`] names the memory a pointer
or tensor lives in;
+//! the functions allocate tensors and query devices and the runtime.
// tvm-ffi-stubgen(missing-docs): allow
// No `prefix`: only the objects with a block below are bound, not all of
`runtime`.
@@ -71,6 +54,10 @@ struct FfiEnumObj {
// tvm-ffi-stubgen(end)
// `ir.PtrType.storage_scope` refers to this type.
+/// A registered memory space, such as `global` or `shared`, that a pointer or
tensor lives in.
+///
+/// Registration order has no meaning, and a backend must support a scope
before it can allocate
+/// storage in it; [`MemoryScope::register`] adds one.
// tvm-ffi-stubgen(begin): object/runtime.MemoryScope
#[allow(missing_docs)]
#[repr(C)]
@@ -114,6 +101,22 @@ impl MemoryScope {
///
/// Returns an error if the allocator of `device` does not support
`mem_scope`; the CPU one supports
/// only the global scope.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::dtype;
+/// use tvm_tirx::runtime;
+/// use tvm_tirx::tvm_ffi::{DLDevice, DLDeviceType};
+///
+/// let cpu = DLDevice::new(DLDeviceType::kDLCPU, 0);
+/// let buffer = runtime::empty(&[2, 8], dtype::float(32), cpu, None)?;
+///
+/// let bytes = runtime::create_view(&buffer, &[64], dtype::int(8), 0)?;
+/// assert_eq!(bytes.shape(), &[64]);
+///
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
pub fn empty(
shape: &[i64],
dtype: DLDataType,
diff --git a/rust/src/script.rs b/rust/src/script.rs
index 7a9e6046f..241703f28 100644
--- a/rust/src/script.rs
+++ b/rust/src/script.rs
@@ -17,28 +17,8 @@
* under the License.
*/
-//! TVMScript: IR objects printed as the Python script that builds them.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::{
-//! DictAttrs, Expr, PrimExpr, PrimType, PrimVar, Return, SeqStmt, Var,
UNKNOWN_LOC,
-//! };
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::tirx::Function;
-//! use tvm_tirx::tvm_ffi::Map;
-//!
-//! let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
-//! assert_eq!(script(&PrimExpr::try_from(&x * 2 + 1)?)?, "x * 2 + 1");
-//!
-//! let body = SeqStmt::new([Return::new(Expr::try_from(&x + 1)?,
UNKNOWN_LOC)], UNKNOWN_LOC);
-//! let attrs = DictAttrs::new(Map::new());
-//! let func = Function::new([&x], &body, PrimType::int(32), attrs,
UNKNOWN_LOC);
-//! let text = script(&func)?;
-//! assert!(text.contains("def main(x: T.int32) -> T.int32:\n return x +
1"));
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! TVMScript printing. [`script`] prints an IR object as TVMScript, the
Python-syntax text form
+//! of the IR.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -59,6 +39,27 @@ use tvm_ffi::String;
/// Print the IR object `obj` as TVMScript, with the default printer config
/// (`script.printer.Script`).
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::{
+/// DictAttrs, Expr, PrimExpr, PrimType, PrimVar, Return, SeqStmt, Var,
UNKNOWN_LOC,
+/// };
+/// use tvm_tirx::script::script;
+/// use tvm_tirx::tirx::Function;
+/// use tvm_tirx::tvm_ffi::Map;
+///
+/// let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
+/// assert_eq!(script(&PrimExpr::try_from(&x * 2 + 1)?)?, "x * 2 + 1");
+///
+/// let body = SeqStmt::new([Return::new(Expr::try_from(&x + 1)?,
UNKNOWN_LOC)], UNKNOWN_LOC);
+/// let attrs = DictAttrs::new(Map::new());
+/// let func = Function::new([&x], &body, PrimType::int(32), attrs,
UNKNOWN_LOC);
+/// let text = script(&func)?;
+/// assert!(text.contains("def main(x: T.int32) -> T.int32:\n return x +
1"));
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
pub fn script<T: ObjectRefCore + AnyCompatible>(obj: &T) -> Result<String> {
tvm_ffi::cached_global_func!("script.printer.Script")
.call_packed(&[AnyView::from(obj), AnyView::new()])?
diff --git a/rust/src/sym.rs b/rust/src/sym.rs
index 64a7a29f2..0931f604f 100644
--- a/rust/src/sym.rs
+++ b/rust/src/sym.rs
@@ -17,10 +17,9 @@
* under the License.
*/
-//! Symbolic analysis of integer expressions.
-//!
-//! An [`Analyzer`] proves and simplifies under the facts of a
[`ConstraintContext`]; [`int_set`]
-//! bounds expressions, and [`iter_affine_map`] detects iterator maps.
+//! Symbolic integer analysis. An [`Analyzer`] proves, bounds and simplifies
integer expressions
+//! under the facts of a [`ConstraintContext`]; [`int_set`] and
[`iter_affine_map`] describe the
+//! values and the iterator structure of index expressions.
// tvm-ffi-stubgen(missing-docs): allow
// An object registered under `sym` that no file has a block for yet is
generated here, to be moved
diff --git a/rust/src/sym/analyzer.rs b/rust/src/sym/analyzer.rs
index 3a9df9b73..e36154674 100644
--- a/rust/src/sym/analyzer.rs
+++ b/rust/src/sym/analyzer.rs
@@ -17,31 +17,9 @@
* under the License.
*/
-//! The analyzer and its results.
-//!
-//! An [`Analyzer`] proves and simplifies expressions under the facts it
knows. [`ConstIntBound`]
-//! and [`ModularSet`] are its bound and modular results.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::prim::op::floordiv;
-//! use tvm_tirx::ir::{PrimType, PrimVar, Range, Var, UNKNOWN_LOC};
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::sym::{Analyzer, ConstraintContext, RangeConstraint};
-//!
-//! let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
-//! let context = ConstraintContext::new();
-//! let domain = Range::from_min_extent(0, 16, UNKNOWN_LOC)?;
-//! context.append(&RangeConstraint::new(x.clone(), domain).into())?;
-//! let analyzer = Analyzer::new();
-//! analyzer.sync_constraint_context(&context)?;
-//!
-//! assert_eq!(script(&analyzer.simplify(floordiv(&x, 16)?, 2)?)?, "0");
-//!
-//! assert_eq!(script(&Analyzer::new().simplify(floordiv(&x, 16)?, 2)?)?, "x
// 16");
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! The analyzer. An [`Analyzer`] proves, bounds and simplifies integer
expressions under the facts
+//! of a [`ConstraintContext`]; its bounds are [`ConstIntBound`]s and
[`ModularSet`]s, and
+//! [`SplitExpr`] and [`SumExpr`] are the forms its canonical simplifier works
in.
// tvm-ffi-stubgen(missing-docs): allow
// Python has no class for `CanonicalExpr`, `SumExpr` and `SplitExpr`, the
forms of the canonical
@@ -94,6 +72,34 @@ use tvm_ffi::VALUE_ERROR;
//
---------------------------------------------------------------------------------------------
// The analyzer
+/// An analyzer of integer expressions: it proves, bounds and simplifies them
under the facts it
+/// knows, which
[`sync_constraint_context`](Analyzer::sync_constraint_context) sets from a
+/// [`ConstraintContext`].
+///
+/// A clone shares the analyzer and its facts;
[`deep_clone`](Analyzer::deep_clone) copies them.
+/// When the build includes the Z3 prover, default-strength proofs fall back
to it unless
+/// [`set_z3_default_fallback`](Analyzer::set_z3_default_fallback) turns that
off.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::prim::op::floordiv;
+/// use tvm_tirx::ir::{PrimType, PrimVar, Range, Var, UNKNOWN_LOC};
+/// use tvm_tirx::script::script;
+/// use tvm_tirx::sym::{Analyzer, ConstraintContext, RangeConstraint};
+///
+/// let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
+/// let context = ConstraintContext::new();
+/// let domain = Range::from_min_extent(0, 16, UNKNOWN_LOC)?;
+/// context.append(&RangeConstraint::new(x.clone(), domain).into())?;
+/// let analyzer = Analyzer::new();
+/// analyzer.sync_constraint_context(&context)?;
+///
+/// assert_eq!(script(&analyzer.simplify(floordiv(&x, 16)?, 2)?)?, "0");
+///
+/// assert_eq!(script(&Analyzer::new().simplify(floordiv(&x, 16)?, 2)?)?, "x
// 16");
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/sym.Analyzer
#[allow(missing_docs)]
#[repr(C)]
@@ -313,7 +319,7 @@ impl Analyzer {
}
}
-/// How hard [`Analyzer::can_prove`] tries, the C++ `sym::ProofStrength`.
+/// How hard [`Analyzer::can_prove`] tries.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
#[repr(i32)]
pub enum ProofStrength {
@@ -324,7 +330,7 @@ pub enum ProofStrength {
SymbolicBound = 1,
}
-/// The result of [`Analyzer::try_compare`], the C++ `sym::CompareResult`.
+/// The result of [`Analyzer::try_compare`].
///
/// The relations other than `Inconsistent` and `Unknown` say what is known of
`lhs` against `rhs`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
@@ -395,6 +401,8 @@ pub fn allow_uint_as_index() ->
Result<AllowUintAsIndexScope> {
//
---------------------------------------------------------------------------------------------
// Results of the analyses
+/// The inclusive constant bounds `[min_value, max_value]` of an integer
expression, where
+/// `i64::MAX` and `-i64::MAX` stand for unbounded ends.
// tvm-ffi-stubgen(begin): object/sym.ConstIntBound
#[allow(missing_docs)]
#[repr(C)]
@@ -442,6 +450,8 @@ impl ConstIntBound {
unsafe impl tvm_ffi::object::ImmutableObject for ConstIntBoundObj {}
// tvm-ffi-stubgen(end)
+/// The congruence class `{coeff * x + base}` over the integers `x` that an
integer expression lies
+/// in: a `coeff` of 0 is the single value `base`, and a `coeff` of 1 any
integer.
// tvm-ffi-stubgen(begin): object/sym.ModularSet
#[allow(missing_docs)]
#[repr(C)]
@@ -489,6 +499,9 @@ impl ModularSet {
unsafe impl tvm_ffi::object::ImmutableObject for ModularSetObj {}
// tvm-ffi-stubgen(end)
+/// The counters of an analyzer's rewrite simplifier: the nodes it visited,
the constraints it
+/// entered, the rewrites it attempted and performed, its deepest recursion
and its recursive
+/// rewrites.
// tvm-ffi-stubgen(begin): object/sym.RewriteSimplifierStats
#[allow(missing_docs)]
#[repr(C)]
@@ -572,6 +585,8 @@ unsafe impl tvm_ffi::object::ImmutableObject for
RewriteSimplifierStatsObj {}
//
---------------------------------------------------------------------------------------------
// The forms of the canonical simplifier
+/// The base of the intermediate forms the canonical simplifier introduces,
[`SplitExpr`] and
+/// [`SumExpr`].
// tvm-ffi-stubgen(begin): object/sym.CanonicalExpr
#[allow(missing_docs)]
#[repr(C)]
@@ -601,6 +616,7 @@ impl Deref for CanonicalExprObj {
tvm_ffi::unsafe_impl_object_ref_upcast!(CanonicalExpr => Expr);
// tvm-ffi-stubgen(end)
+/// The canonical form of a split index: `((index % upper_factor) /
lower_factor) * scale`.
// tvm-ffi-stubgen(begin): object/sym.SplitExpr
#[allow(missing_docs)]
#[repr(C)]
@@ -630,6 +646,7 @@ impl Deref for SplitExprObj {
tvm_ffi::unsafe_impl_object_ref_upcast!(SplitExpr => Expr, SplitExpr =>
CanonicalExpr, SplitExpr => PrimExpr);
// tvm-ffi-stubgen(end)
+/// The canonical form of a sum: a sum of [`SplitExpr`]s plus a constant base.
// tvm-ffi-stubgen(begin): object/sym.SumExpr
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/sym/constraint.rs b/rust/src/sym/constraint.rs
index a41d2d408..4bcff25dd 100644
--- a/rust/src/sym/constraint.rs
+++ b/rust/src/sym/constraint.rs
@@ -19,32 +19,9 @@
// The generated `new` of a field-less base needs no `Default`.
#![allow(clippy::new_without_default)]
-//! Constraints: the facts in scope at a point of a traversal.
-//!
-//! A [`ConstraintContext`] holds [`BindingConstraint`]s, [`RangeConstraint`]s
and
-//! [`PredicateConstraint`]s, ancestors' first. Its handles are neither `Send`
nor `Sync`; do not
-//! share a context across threads, also not through an erased `ObjectRef`.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::prim::op::{ge, lt};
-//! use tvm_tirx::ir::{PrimType, PrimVar, Var, UNKNOWN_LOC};
-//! use tvm_tirx::sym::{Analyzer, ConstraintContext, PredicateConstraint,
ProofStrength};
-//!
-//! let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
-//! let outer = ConstraintContext::new();
-//! outer.append(&PredicateConstraint::new(ge(&x, 0)?, false).into())?;
-//! let inner = outer.snapshot()?;
-//! inner.append(&PredicateConstraint::new(lt(&x, 16)?, false).into())?;
-//!
-//! let analyzer = Analyzer::new();
-//! analyzer.sync_constraint_context(&inner)?;
-//! assert!(analyzer.can_prove(lt(&x, 16)?, ProofStrength::Default)?);
-//! analyzer.sync_constraint_context(&outer)?;
-//! assert!(!analyzer.can_prove(lt(&x, 16)?, ProofStrength::Default)?);
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Constraints. A [`Constraint`] is one fact in scope, a
[`BindingConstraint`],
+//! [`RangeConstraint`] or [`PredicateConstraint`]; a [`ConstraintContext`]
holds the facts in
+//! scope at a point of a traversal, which an [`Analyzer`](super::Analyzer)
proves under.
// tvm-ffi-stubgen(missing-docs): allow
// tvm-ffi-stubgen(ty-map): sym.ConstraintContext ->
crate::sym::constraint::ConstraintContext
@@ -68,6 +45,7 @@ use tvm_ffi::ShallowCopy;
use tvm_ffi::UnsafeInit;
// tvm-ffi-stubgen(end)
+/// The base of every fact a [`ConstraintContext`] holds; it owns no analyzer
state.
// tvm-ffi-stubgen(begin): object/sym.Constraint
#[allow(missing_docs)]
#[repr(C)]
@@ -126,6 +104,8 @@ unsafe impl UnsafeInit for ConstraintObj {
unsafe impl tvm_ffi::object::ImmutableObject for ConstraintObj {}
// tvm-ffi-stubgen(end)
+/// The binding of `var` to `value`. The value is the bound expression as
written, memory reads
+/// included: it records what `var` was bound to, not an equality that holds
afterwards.
// tvm-ffi-stubgen(begin): object/sym.BindingConstraint
#[allow(missing_docs)]
#[repr(C)]
@@ -182,6 +162,8 @@ unsafe impl tvm_ffi::object::ImmutableObject for
BindingConstraintObj {}
tvm_ffi::unsafe_impl_object_ref_upcast!(BindingConstraint => Constraint);
// tvm-ffi-stubgen(end)
+/// The fact that the boolean `condition` holds. With `assume_always`, the
caller guarantees it
+/// holds throughout the scope, even after the values it reads change.
// tvm-ffi-stubgen(begin): object/sym.PredicateConstraint
#[allow(missing_docs)]
#[repr(C)]
@@ -238,6 +220,8 @@ unsafe impl tvm_ffi::object::ImmutableObject for
PredicateConstraintObj {}
tvm_ffi::unsafe_impl_object_ref_upcast!(PredicateConstraint => Constraint);
// tvm-ffi-stubgen(end)
+/// The fact that `var` lies in the half-open `range`, which may be wider than
the values `var`
+/// takes.
// tvm-ffi-stubgen(begin): object/sym.RangeConstraint
#[allow(missing_docs)]
#[repr(C)]
@@ -323,7 +307,29 @@ const _: () = {
/// A context of constraints: the premises in effect at a point of a
traversal, ancestors' first.
///
/// Clones share the context's identity and later appends; use
-/// [`snapshot`](ConstraintContextObj::snapshot) to capture the current prefix.
+/// [`snapshot`](ConstraintContextObj::snapshot) to capture the current
prefix. A context is neither
+/// `Send` nor `Sync`: do not share one across threads, also not through an
erased `ObjectRef`.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::prim::op::{ge, lt};
+/// use tvm_tirx::ir::{PrimType, PrimVar, Var, UNKNOWN_LOC};
+/// use tvm_tirx::sym::{Analyzer, ConstraintContext, PredicateConstraint,
ProofStrength};
+///
+/// let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
+/// let outer = ConstraintContext::new();
+/// outer.append(&PredicateConstraint::new(ge(&x, 0)?, false).into())?;
+/// let inner = outer.snapshot()?;
+/// inner.append(&PredicateConstraint::new(lt(&x, 16)?, false).into())?;
+///
+/// let analyzer = Analyzer::new();
+/// analyzer.sync_constraint_context(&inner)?;
+/// assert!(analyzer.can_prove(lt(&x, 16)?, ProofStrength::Default)?);
+/// analyzer.sync_constraint_context(&outer)?;
+/// assert!(!analyzer.can_prove(lt(&x, 16)?, ProofStrength::Default)?);
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
pub struct ConstraintContext {
@@ -393,8 +399,7 @@ impl ConstraintContextObj {
}
impl ConstraintContext {
- /// Create an empty root context, with the complete native-compatible
layout (C++
- /// `ConstraintContext()`).
+ /// Create an empty root context, with the complete native-compatible
layout.
pub fn new() -> Self {
Self::allocate(None, 0)
}
diff --git a/rust/src/sym/int_set.rs b/rust/src/sym/int_set.rs
index 52dc5a9a5..014296323 100644
--- a/rust/src/sym/int_set.rs
+++ b/rust/src/sym/int_set.rs
@@ -17,28 +17,9 @@
* under the License.
*/
-//! Sets of integers: [`IntSet`] and the interval [`IntervalSet`], with bound
deduction and region
-//! bound estimation.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::prim::op::PrimOps;
-//! use tvm_tirx::ir::{PrimType, PrimVar, Var, UNKNOWN_LOC};
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::sym::{self, Analyzer, IntSet, IntervalSet};
-//! use tvm_tirx::tvm_ffi::Map;
-//!
-//! let i = PrimVar::new("i", PrimType::int(32), UNKNOWN_LOC);
-//! let dom_map = Map::from_iter([(Var::from(&i),
IntSet::from(IntervalSet::new(0, 15)))]);
-//! let set = Analyzer::new().int_set(&i * 2 + 1, Some(&dom_map))?;
-//! assert_eq!(script(&set)?, "sym.IntervalSet(min_value=1, max_value=31)");
-//!
-//! let none = Map::<Var, IntSet>::new();
-//! let bound = sym::deduce_bound(&i, (&i * 4 + 3).less(64), &none, &none)?;
-//! assert_eq!(script(&bound)?, "sym.IntervalSet(min_value=neg_inf,
max_value=15)");
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Integer sets. An [`IntSet`], such as the interval [`IntervalSet`], bounds
the values an
+//! expression takes; [`deduce_bound`] solves a condition for a variable, and
the
+//! `estimate_region_*` functions bound the region a set of accesses covers.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -73,6 +54,28 @@ use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
// tvm-ffi-stubgen(end)
+/// The base of every set of integers, one dimension each, that bounds the
values an expression
+/// takes; [`Analyzer::int_set`] computes one.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::prim::op::PrimOps;
+/// use tvm_tirx::ir::{PrimType, PrimVar, Var, UNKNOWN_LOC};
+/// use tvm_tirx::script::script;
+/// use tvm_tirx::sym::{self, Analyzer, IntSet, IntervalSet};
+/// use tvm_tirx::tvm_ffi::Map;
+///
+/// let i = PrimVar::new("i", PrimType::int(32), UNKNOWN_LOC);
+/// let dom_map = Map::from_iter([(Var::from(&i),
IntSet::from(IntervalSet::new(0, 15)))]);
+/// let set = Analyzer::new().int_set(&i * 2 + 1, Some(&dom_map))?;
+/// assert_eq!(script(&set)?, "sym.IntervalSet(min_value=1, max_value=31)");
+///
+/// let none = Map::<Var, IntSet>::new();
+/// let bound = sym::deduce_bound(&i, (&i * 4 + 3).less(64), &none, &none)?;
+/// assert_eq!(script(&bound)?, "sym.IntervalSet(min_value=neg_inf,
max_value=15)");
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/sym.IntSet
#[allow(missing_docs)]
#[repr(C)]
@@ -120,6 +123,8 @@ impl IntSet {
}
}
+/// The interval `[min_value, max_value]` of integers, whose bounds are
symbolic expressions and
+/// may be [`neg_inf`] or [`pos_inf`].
// tvm-ffi-stubgen(begin): object/sym.IntervalSet
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/sym/iter_affine_map.rs b/rust/src/sym/iter_affine_map.rs
index 891af39b1..278140833 100644
--- a/rust/src/sym/iter_affine_map.rs
+++ b/rust/src/sym/iter_affine_map.rs
@@ -17,34 +17,9 @@
* under the License.
*/
-//! Iterator maps: index expressions as sums of splits of iterators.
-//!
-//! An [`IterSumExpr`] is a sum of [`IterSplitExpr`]s `(source //
lower_factor) % extent * scale`
-//! plus `base_`. [`detect_iter_map`] detects them, and its [`IterMapResult`]
reports `errors`.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::{PrimType, PrimVar, Range, Var, UNKNOWN_LOC};
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::sym::{self, IterMapLevel};
-//! use tvm_tirx::tvm_ffi::Map;
-//!
-//! let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
-//! let y = PrimVar::new("y", PrimType::int(32), UNKNOWN_LOC);
-//! let four = Range::from_min_extent(0, 4, UNKNOWN_LOC)?;
-//! let iters = Map::from_iter([(x.clone(), four.clone()), (y.clone(), four)]);
-//!
-//! let bijective = IterMapLevel::Bijective;
-//! let map = sym::detect_iter_map([&x * 4 + &y], &iters, true, bijective,
None, true)?;
-//! assert!(map.errors.is_empty());
-//!
-//! let o = PrimVar::new("o", PrimType::int(32), UNKNOWN_LOC);
-//! let inverse = sym::inverse_affine_iter_map(map.indices.iter(), [&o])?;
-//! assert_eq!(script(&inverse.get(&Var::from(&x))?.unwrap())?, "o // 4");
-//! assert_eq!(script(&inverse.get(&Var::from(&y))?.unwrap())?, "o % 4");
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Iterator maps. An iterator map writes index expressions as
[`IterSumExpr`]s, sums of
+//! [`IterSplitExpr`]s of [`IterMark`]ed iterators; [`detect_iter_map`] finds
one, and the other
+//! functions simplify, divide and invert such maps.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -95,6 +70,8 @@ use tvm_ffi::Result;
use tvm_ffi::String;
// tvm-ffi-stubgen(end)
+/// The base of the iterator map forms, [`IterSplitExpr`] and [`IterSumExpr`].
They hold the result
+/// of detection and do not appear in a well-formed function.
// tvm-ffi-stubgen(begin): object/sym.IterMapExpr
#[allow(missing_docs)]
#[repr(C)]
@@ -124,6 +101,36 @@ impl Deref for IterMapExprObj {
tvm_ffi::unsafe_impl_object_ref_upcast!(IterMapExpr => Expr, IterMapExpr =>
PrimExpr);
// tvm-ffi-stubgen(end)
+/// The result of [`detect_iter_map`]: the detected `indices`, or empty
`indices` and the `errors`
+/// that say why detection failed.
+///
+/// `padding_predicate`, over the input iterators, holds for the indices that
the returned
+/// iterators cover beyond the bounds of the input ones; it is `None` if
detection failed before
+/// computing it.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::{PrimType, PrimVar, Range, Var, UNKNOWN_LOC};
+/// use tvm_tirx::script::script;
+/// use tvm_tirx::sym::{self, IterMapLevel};
+/// use tvm_tirx::tvm_ffi::Map;
+///
+/// let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
+/// let y = PrimVar::new("y", PrimType::int(32), UNKNOWN_LOC);
+/// let four = Range::from_min_extent(0, 4, UNKNOWN_LOC)?;
+/// let iters = Map::from_iter([(x.clone(), four.clone()), (y.clone(), four)]);
+///
+/// let bijective = IterMapLevel::Bijective;
+/// let map = sym::detect_iter_map([&x * 4 + &y], &iters, true, bijective,
None, true)?;
+/// assert!(map.errors.is_empty());
+///
+/// let o = PrimVar::new("o", PrimType::int(32), UNKNOWN_LOC);
+/// let inverse = sym::inverse_affine_iter_map(map.indices.iter(), [&o])?;
+/// assert_eq!(script(&inverse.get(&Var::from(&x))?.unwrap())?, "o // 4");
+/// assert_eq!(script(&inverse.get(&Var::from(&y))?.unwrap())?, "o % 4");
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/sym.IterMapResult
#[allow(missing_docs)]
#[repr(C)]
@@ -178,6 +185,7 @@ impl IterMapResult {
}
// tvm-ffi-stubgen(end)
+/// The mark of `source` as an iterator over `[0, extent)`.
// tvm-ffi-stubgen(begin): object/sym.IterMark
#[allow(missing_docs)]
#[repr(C)]
@@ -225,6 +233,7 @@ impl IterMark {
unsafe impl tvm_ffi::object::ImmutableObject for IterMarkObj {}
// tvm-ffi-stubgen(end)
+/// The split `(source // lower_factor) % extent * scale` of the marked
iterator `source`.
// tvm-ffi-stubgen(begin): object/sym.IterSplitExpr
#[allow(missing_docs)]
#[repr(C)]
@@ -276,6 +285,7 @@ unsafe impl tvm_ffi::object::ImmutableObject for
IterSplitExprObj {}
tvm_ffi::unsafe_impl_object_ref_upcast!(IterSplitExpr => Expr, IterSplitExpr
=> IterMapExpr, IterSplitExpr => PrimExpr);
// tvm-ffi-stubgen(end)
+/// The sum of the splits `args` plus `base`.
// tvm-ffi-stubgen(begin): object/sym.IterSumExpr
#[allow(missing_docs)]
#[repr(C)]
@@ -355,7 +365,7 @@ impl IterSumExpr {
}
}
-/// How strictly an iterator map is checked, the C++ `sym::IterMapLevel`.
+/// How strictly an iterator map is checked.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
#[repr(i32)]
pub enum IterMapLevel {
diff --git a/rust/src/target.rs b/rust/src/target.rs
index 6b76b9b86..5ccd51896 100644
--- a/rust/src/target.rs
+++ b/rust/src/target.rs
@@ -17,26 +17,9 @@
* under the License.
*/
-//! Compilation targets.
-//!
-//! A [`Target`] has a [`TargetKind`], keys, attributes and an optional host.
It is created from a
-//! tag, a kind name or a config, and made the current target with
[`Target::scope`].
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::target::Target;
-//! use tvm_tirx::tvm_ffi::ObjectRefCore;
-//!
-//! let cpu = Target::new(r#"{"kind": "llvm", "mcpu": "skylake"}"#)?;
-//! assert!(cpu.as_json()?.as_str().contains(r#""keys":["cpu"]"#));
-//!
-//! let scope = cpu.scope()?;
-//! assert!(Target::current(false)?.unwrap().same_as(&cpu));
-//! drop(scope);
-//! assert!(Target::current(true)?.is_none());
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Compilation targets. A [`Target`] describes the device and options code is
compiled for, as an
+//! instance of a registered [`TargetKind`]; it is created from a tag, a kind
name or a config, and
+//! [`Target::scope`] makes it the current one.
// tvm-ffi-stubgen(missing-docs): allow
// Every object registered under `target` has its block in this file.
@@ -61,6 +44,25 @@ use tvm_ffi::String;
use tvm_ffi::object::ObjectRef;
// tvm-ffi-stubgen(end)
+/// A compilation target: its `kind`, the `keys` that classify it, such as
`cpu`, the `attrs` of
+/// its options, and an optional `host` target for the code that runs on the
host. A target may come
+/// from a registered `tag`, such as `nvidia/nvidia-a100`.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::target::Target;
+/// use tvm_tirx::tvm_ffi::ObjectRefCore;
+///
+/// let cpu = Target::new(r#"{"kind": "llvm", "mcpu": "skylake"}"#)?;
+/// assert!(cpu.as_json()?.as_str().contains(r#""keys":["cpu"]"#));
+///
+/// let scope = cpu.scope()?;
+/// assert!(Target::current(false)?.unwrap().same_as(&cpu));
+/// drop(scope);
+/// assert!(Target::current(true)?.is_none());
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/target.Target
#[allow(missing_docs)]
#[repr(C)]
@@ -201,6 +203,8 @@ impl Drop for TargetScope {
}
}
+/// A kind of target, such as `llvm` or `cuda`: its `name`, the options it
takes, and the device
+/// type and keys of its targets by default.
// tvm-ffi-stubgen(begin): object/target.TargetKind
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/tirx.rs b/rust/src/tirx.rs
index 884fe4c15..d0607d711 100644
--- a/rust/src/tirx.rs
+++ b/rust/src/tirx.rs
@@ -17,10 +17,9 @@
* under the License.
*/
-//! Tensor programs: functions, tensor types, index maps, layouts and thread
scopes.
-//!
-//! [`structural`] holds the traversals that keep the facts in scope, [`tile`]
the tile
-//! instructions, [`transform`] the passes and [`analysis`] the analyses.
+//! Tensor programs. A [`Function`] computes on tensors, variables typed by a
[`TensorType`] whose
+//! [`Layout`] maps their indices to memory and threads; [`tile`] instructions
are dispatched per
+//! [`ThreadScope`], and [`transform`] and [`analysis`] hold the passes and
analyses.
// tvm-ffi-stubgen(missing-docs): allow
// An object registered under `tirx` that no file has a block for yet is
generated here, to be moved
diff --git a/rust/src/tirx/analysis.rs b/rust/src/tirx/analysis.rs
index 41916e22d..9ea3bd6c3 100644
--- a/rust/src/tirx/analysis.rs
+++ b/rust/src/tirx/analysis.rs
@@ -17,8 +17,8 @@
* under the License.
*/
-//! Analyses of tensor-level IR: memory accesses, undefined variables, linear
patterns and
-//! well-formedness.
+//! Analyses. Read-only queries over functions, statements and expressions:
well-formedness and
+//! host memory access checks, undefined variables, and linear and clip-bound
index patterns.
#[allow(clippy::module_inception)]
pub mod analysis;
diff --git a/rust/src/tirx/analysis/analysis.rs
b/rust/src/tirx/analysis/analysis.rs
index 3c6900bce..445c7c990 100644
--- a/rust/src/tirx/analysis/analysis.rs
+++ b/rust/src/tirx/analysis/analysis.rs
@@ -17,22 +17,7 @@
* under the License.
*/
-//! The analyses of tensor-level IR (`tirx.analysis.*`).
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::{Expr, PrimType, PrimVar, Var, UNKNOWN_LOC};
-//! use tvm_tirx::tirx::analysis;
-//! use tvm_tirx::tvm_ffi::ObjectRefCore;
-//!
-//! let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
-//! let y = PrimVar::new("y", PrimType::int(32), UNKNOWN_LOC);
-//! let sum: Expr = (&x + &y).try_into()?;
-//! let undefined = analysis::undefined_vars(sum, [&x])?;
-//! assert!(undefined.len() == 1 && undefined.get(0)?.same_as(&y));
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! The analyses, which [`analysis`](super) re-exports.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -76,6 +61,21 @@ pub fn verify_memory(func: impl IntoObjectRef<Function>) ->
Result<bool> {
/// Return the variables `node`, a statement or an expression, uses but
neither defines nor finds in
/// `defs` (`tirx.analysis.UndefinedVars`).
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::{Expr, PrimType, PrimVar, Var, UNKNOWN_LOC};
+/// use tvm_tirx::tirx::analysis;
+/// use tvm_tirx::tvm_ffi::ObjectRefCore;
+///
+/// let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
+/// let y = PrimVar::new("y", PrimType::int(32), UNKNOWN_LOC);
+/// let sum: Expr = (&x + &y).try_into()?;
+/// let undefined = analysis::undefined_vars(sum, [&x])?;
+/// assert!(undefined.len() == 1 && undefined.get(0)?.same_as(&y));
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
pub fn undefined_vars(
node: impl Into<Any>,
defs: impl IntoIterator<Item = impl Into<Var>>,
diff --git a/rust/src/tirx/function.rs b/rust/src/tirx/function.rs
index b550375cc..f8a31105a 100644
--- a/rust/src/tirx/function.rs
+++ b/rust/src/tirx/function.rs
@@ -17,32 +17,13 @@
* under the License.
*/
-//! Tensor functions.
-//!
-//! A [`Function`] is a [`BaseFunc`] with parameters, a return type and a
sequence body.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::{DictAttrs, PrimType, PrimVar, Return, SeqStmt, Var,
UNKNOWN_LOC};
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::tirx::{renew_def, Function};
-//! use tvm_tirx::tvm_ffi::{Map, ObjectRefCore};
-//!
-//! let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
-//! let body = SeqStmt::new([Return::new(&x, UNKNOWN_LOC)], UNKNOWN_LOC);
-//! let attrs = DictAttrs::new(Map::new());
-//! let identity = Function::new([&x], &body, PrimType::int(32), attrs,
UNKNOWN_LOC);
-//! assert!(script(&identity)?.as_str().contains("def main(x: T.int32) ->
T.int32:"));
-//!
-//! let copy = renew_def(&identity)?;
-//! assert!(!copy.params.get(0)?.same_as(&x));
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Functions of the dialect. A [`Function`] is a [`BaseFunc`] whose type is
derived from its
+//! parameters and return type, with a [`SeqStmt`] body; [`renew_def`] copies
one with fresh
+//! variables.
// tvm-ffi-stubgen(missing-docs): allow
-// `custom-new` reserves `new` for a constructor with the C++ one's semantics
(`ty` derived from the
-// signature); the generated `from_fields` allocates the fields as given.
+// `custom-new` reserves `new` for a hand-written constructor (`ty` derived
from the signature); the
+// generated `from_fields` allocates the fields as given.
// tvm-ffi-stubgen(custom-new): tirx.Function
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -74,6 +55,27 @@ use tvm_ffi::ObjectRefCore;
use tvm_ffi::Result;
// tvm-ffi-stubgen(end)
+/// A function: `params`, a `body` and a `ret_type`, typed by the [`FuncType`]
of its parameter
+/// and return types. A declaration has an empty body and the `declaration`
attribute set.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::{DictAttrs, PrimType, PrimVar, Return, SeqStmt, Var,
UNKNOWN_LOC};
+/// use tvm_tirx::script::script;
+/// use tvm_tirx::tirx::{renew_def, Function};
+/// use tvm_tirx::tvm_ffi::{Map, ObjectRefCore};
+///
+/// let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
+/// let body = SeqStmt::new([Return::new(&x, UNKNOWN_LOC)], UNKNOWN_LOC);
+/// let attrs = DictAttrs::new(Map::new());
+/// let identity = Function::new([&x], &body, PrimType::int(32), attrs,
UNKNOWN_LOC);
+/// assert!(script(&identity)?.as_str().contains("def main(x: T.int32) ->
T.int32:"));
+///
+/// let copy = renew_def(&identity)?;
+/// assert!(!copy.params.get(0)?.same_as(&x));
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/tirx.Function
#[allow(missing_docs)]
#[repr(C)]
@@ -153,8 +155,8 @@ tvm_ffi::unsafe_impl_object_ref_upcast!(Function => Expr,
Function => BaseFunc);
// tvm-ffi-stubgen(end)
impl Function {
- /// Create the function of `params`, `body`, `ret_type` and `attrs`; a
missing return type
- /// becomes void (C++ `Function(params, body, ret_type, attrs, loc)`).
+ /// Create the function of `params`, `body`, `ret_type` and `attrs`, typed
by the function
+ /// type of its signature; a missing return type becomes void, the empty
tuple type.
pub fn new(
params: impl IntoIterator<Item = impl Into<Var>>,
body: &SeqStmt,
diff --git a/rust/src/tirx/index_map.rs b/rust/src/tirx/index_map.rs
index 8354a78d2..714967cbe 100644
--- a/rust/src/tirx/index_map.rs
+++ b/rust/src/tirx/index_map.rs
@@ -17,29 +17,9 @@
* under the License.
*/
-//! Index maps.
-//!
-//! An [`IndexMap`] maps its `initial_indices` to its `final_indices`, as `(i,
j) -> (j, i)` does.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::{PrimExpr, PrimType, PrimVar, UNKNOWN_LOC};
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::tirx::IndexMap;
-//! use tvm_tirx::tvm_ffi::{Array, Result};
-//!
-//! let i = PrimVar::new("i", PrimType::int(32), UNKNOWN_LOC);
-//! let j = PrimVar::new("j", PrimType::int(32), UNKNOWN_LOC);
-//! let transpose = IndexMap::new([&i, &j], [&j, &i], None)?;
-//! assert_eq!(script(&transpose)?, "T.index_map(lambda i, j: (j, i))");
-//!
-//! let text =
-//! |exprs: Array<PrimExpr>|
exprs.iter().map(script).collect::<Result<Vec<_>>>();
-//! assert_eq!(text(transpose.map_indices([1, 2], None)?)?, ["2", "1"]);
-//!
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Index maps. An [`IndexMap`] maps the indices of one representation of a
tensor to those of
+//! another, as a layout transformation does; it maps indices, shapes and
tensor data, and inverts
+//! when it is bijective.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -51,7 +31,7 @@
// tvm-ffi-stubgen(import-object): tvm_ffi::Tensor
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimExpr
-// `custom-new` reserves `new` for the C++ constructor, which takes the
inverse as an `IndexMap`; the
+// `custom-new` reserves `new` for a constructor that takes the inverse as an
`IndexMap`; the
// generated `from_fields` stores it as the `ObjectRef` the field holds.
// tvm-ffi-stubgen(custom-new): tirx.IndexMap
// The indices a map computes are `PrimExpr`s, as C++ `IndexMapObj` holds them.
@@ -75,6 +55,28 @@ use tvm_ffi::Tensor;
use tvm_ffi::object::ObjectRef;
// tvm-ffi-stubgen(end)
+/// The map from the index variables `initial_indices` to the expressions
`final_indices` over
+/// them, such as `(i, j) -> (j, i)`, with its inverse `inverse_index_map`
when it is known.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::{PrimExpr, PrimType, PrimVar, UNKNOWN_LOC};
+/// use tvm_tirx::script::script;
+/// use tvm_tirx::tirx::IndexMap;
+/// use tvm_tirx::tvm_ffi::{Array, Result};
+///
+/// let i = PrimVar::new("i", PrimType::int(32), UNKNOWN_LOC);
+/// let j = PrimVar::new("j", PrimType::int(32), UNKNOWN_LOC);
+/// let transpose = IndexMap::new([&i, &j], [&j, &i], None)?;
+/// assert_eq!(script(&transpose)?, "T.index_map(lambda i, j: (j, i))");
+///
+/// let text =
+/// |exprs: Array<PrimExpr>|
exprs.iter().map(script).collect::<Result<Vec<_>>>();
+/// assert_eq!(text(transpose.map_indices([1, 2], None)?)?, ["2", "1"]);
+///
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/tirx.IndexMap
#[allow(missing_docs)]
#[repr(C)]
@@ -137,8 +139,7 @@ unsafe impl tvm_ffi::object::ImmutableObject for
IndexMapObj {}
impl IndexMap {
/// Create the map from `initial_indices` to `final_indices`, with the
inverse
- /// `inverse_index_map` if it is known (C++ `IndexMap(initial_indices,
final_indices,
- /// inverse_index_map)`).
+ /// `inverse_index_map` if it is known.
pub fn new(
initial_indices: impl IntoIterator<Item = impl Into<PrimVar>>,
final_indices: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
diff --git a/rust/src/tirx/layout.rs b/rust/src/tirx/layout.rs
index da61e62b6..ba8362e09 100644
--- a/rust/src/tirx/layout.rs
+++ b/rust/src/tirx/layout.rs
@@ -17,32 +17,9 @@
* under the License.
*/
-//! Layouts: how the indices of a tensor map to memory and to threads.
-//!
-//! A [`TileLayout`] maps indices through the [`Iter`]s of its `shard`
(distinct elements) and
-//! `replica` (copies), each an extent and a stride over an [`Axis`]; a
[`ComposeLayout`] swizzles
-//! its memory coordinate. Both are [`Layout`]s.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::{PrimType, PrimVar, Var, UNKNOWN_LOC};
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::sym::Analyzer;
-//! use tvm_tirx::tirx::{Axis, Iter, Layout, TileLayout};
-//! use tvm_tirx::tvm_ffi::Map;
-//!
-//! let m = Axis::get("m");
-//! let shard = [Iter::new(8, 8, m.clone()), Iter::new(8, 1, m)];
-//! let layout = Layout::from(&TileLayout::new(shard, Vec::<Iter>::new(),
&Map::new()));
-//! assert_eq!(script(&layout.size(None)?)?, "64");
-//!
-//! let i = PrimVar::new("i", PrimType::int(32), UNKNOWN_LOC);
-//! let j = PrimVar::new("j", PrimType::int(32), UNKNOWN_LOC);
-//! let offset = layout.apply([&i, &j],
&Analyzer::new())?.get(&"m".into())?.unwrap();
-//! assert_eq!(script(&offset)?, "i * 8 + j");
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Layouts. A [`Layout`] maps the logical indices of a tensor to coordinates
on named [`Axis`]es,
+//! a memory offset and the threads that hold each element; a [`TileLayout`]
builds the map from
+//! strided [`Iter`] factors, and a [`ComposeLayout`] swizzles the memory
coordinate of one.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -89,6 +66,33 @@ struct FfiEnumObj {
}
// tvm-ffi-stubgen(end)
+/// The base of every layout: a map from the logical indices of a tensor to a
coordinate on each
+/// of its [`Axis`]es, such as the memory axis `m` or a thread axis.
+///
+/// [`apply`](Layout::apply) maps one index, [`size`](Layout::size) is the
number of elements the
+/// layout maps and [`span`](Layout::span) the extent of the coordinates it
reaches. Each kind of
+/// layout, such as [`TileLayout`] or [`ComposeLayout`], is a subtype.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::{PrimType, PrimVar, Var, UNKNOWN_LOC};
+/// use tvm_tirx::script::script;
+/// use tvm_tirx::sym::Analyzer;
+/// use tvm_tirx::tirx::{Axis, Iter, Layout, TileLayout};
+/// use tvm_tirx::tvm_ffi::Map;
+///
+/// let m = Axis::get("m");
+/// let shard = [Iter::new(8, 8, m.clone()), Iter::new(8, 1, m)];
+/// let layout = Layout::from(&TileLayout::new(shard, Vec::<Iter>::new(),
&Map::new()));
+/// assert_eq!(script(&layout.size(None)?)?, "64");
+///
+/// let i = PrimVar::new("i", PrimType::int(32), UNKNOWN_LOC);
+/// let j = PrimVar::new("j", PrimType::int(32), UNKNOWN_LOC);
+/// let offset = layout.apply([&i, &j],
&Analyzer::new())?.get(&"m".into())?.unwrap();
+/// assert_eq!(script(&offset)?, "i * 8 + j");
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/tirx.Layout
#[allow(missing_docs)]
#[repr(C)]
@@ -222,6 +226,9 @@ impl Layout {
}
}
+/// The layout of strided factors: each [`Iter`] of `shard` spans distinct
elements and each of
+/// `replica` their copies, adding its stride per position to the coordinate
on its axis, from the
+/// origin `offset` gives on each axis.
// tvm-ffi-stubgen(begin): object/tirx.TileLayout
#[allow(missing_docs)]
#[repr(C)]
@@ -334,6 +341,8 @@ impl TileLayout {
}
}
+/// One factor of a [`TileLayout`]: `extent` positions, each adding `stride`
to the coordinate on
+/// `axis`.
// tvm-ffi-stubgen(begin): object/tirx.Iter
#[allow(missing_docs)]
#[repr(C)]
@@ -388,6 +397,9 @@ unsafe impl tvm_ffi::object::ImmutableObject for IterObj {}
// The native `tirx.Iter` stores its fields as given, as the generated
`Iter::new` does.
+/// A named coordinate a layout maps to: the memory axis `m`, or a thread
axis, such as `laneid`,
+/// that indexes the threads of a scope. Each name is one registered axis,
which [`Axis::get`]
+/// looks up.
// tvm-ffi-stubgen(begin): object/tirx.Axis
#[allow(missing_docs)]
#[repr(C)]
@@ -449,6 +461,9 @@ impl Axis {
}
}
+/// The layout `tile_layout` with its memory coordinate XOR-swizzled: above
the low `per_element`
+/// bits, the group of `swizzle_len` bits `atom_len` bits up is XORed into the
group below if
+/// `swizzle_inner`, and the lower group into the upper one otherwise.
// tvm-ffi-stubgen(begin): object/tirx.ComposeLayout
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/tirx/op.rs b/rust/src/tirx/op.rs
index 76448d949..b5262acd5 100644
--- a/rust/src/tirx/op.rs
+++ b/rust/src/tirx/op.rs
@@ -17,24 +17,9 @@
* under the License.
*/
-//! Attributes of TIRx operator calls.
-//!
-//! The functions build TIRx operator calls; a [`CallFFIKernelAttr`] holds the
launch parameters of
-//! a `call_ffi_kernel`.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::{PrimType, PrimVar, Var, UNKNOWN_LOC};
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::tirx::{logaddexp, reinterpret};
-//!
-//! let x = PrimVar::new("x", PrimType::float(32), UNKNOWN_LOC);
-//! let y = PrimVar::new("y", PrimType::float(32), UNKNOWN_LOC);
-//! assert_eq!(script(&logaddexp(&x, &y)?)?, "T.log(T.exp(x) + T.exp(y))");
-//! assert_eq!(script(&reinterpret(PrimType::int(32), &x)?)?,
"T.reinterpret(x, ty=T.int32)");
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Operators of the dialect. The functions build calls of its operators, such
as [`logaddexp`]
+//! and [`reinterpret`]; a [`CallFFIKernelAttr`] holds the launch metadata of
a `call_ffi_kernel`
+//! call.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -96,6 +81,13 @@ use tvm_ffi::String;
use tvm_ffi::TYPE_ERROR;
// tvm-ffi-stubgen(end)
+/// The launch metadata of a `call_ffi_kernel` call, whose arguments are the
kernel symbol, the
+/// kernel operands and then the launch values.
+///
+/// `launch_params` tags the launch values in order (a flag-only tag takes no
value, and the dynamic
+/// shared-memory size, when present, comes last), `launch_fields` names the
CUDA launch
+/// configuration field each launch value sets, `num_kernel_args` counts the
kernel operands, and
+/// `kernel_attrs` holds the compile-time requirements of the kernel.
// tvm-ffi-stubgen(begin): object/tirx.CallFFIKernelAttr
#[allow(missing_docs)]
#[repr(C)]
@@ -172,8 +164,7 @@ tvm_ffi::unsafe_impl_object_ref_upcast!(CallFFIKernelAttr
=> Attrs);
// The native `tirx.CallFFIKernelAttr` stores its fields as given, as the
generated
// `CallFFIKernelAttr::new` does.
-/// Build `log(exp(a) + exp(b))` of the floating-point `a` and `b`, matched to
one type
-/// (C++ `tirx::logaddexp`, `tirx._OpLogAddExp`).
+/// Build `log(exp(a) + exp(b))` of the floating-point `a` and `b`, matched to
one type.
///
/// # Errors
///
@@ -200,12 +191,23 @@ pub fn logaddexp_at(
log(add(exp(a)?, exp(b)?)?)
}
-/// Reinterpret `value` bit for bit as the primitive or pointer type `ty`
-/// (C++ `tirx::reinterpret(target_ty, value, loc)`, `tirx.reinterpret`).
+/// Reinterpret `value` bit for bit as the primitive or pointer type `ty`
(`tirx.reinterpret`).
///
/// Besides primitive types of the same size, a string reinterprets as a
pointer and a pointer as a
/// scalar `int64`, and back.
///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::{PrimType, PrimVar, UNKNOWN_LOC};
+/// use tvm_tirx::script::script;
+/// use tvm_tirx::tirx::reinterpret;
+///
+/// let x = PrimVar::new("x", PrimType::float(32), UNKNOWN_LOC);
+/// assert_eq!(script(&reinterpret(PrimType::int(32), &x)?)?,
"T.reinterpret(x, ty=T.int32)");
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
+///
/// # Errors
///
/// Returns an error for any other pair of types, or two primitive types of
different sizes.
@@ -307,10 +309,10 @@ pub fn reinterpret_at(
}
/// Build a rational approximation of `erf(arg)` on `[-4, 4]` in `bits`-bit
floats, as the Metal and
-/// WebGPU backends lower `erf` (C++ `tirx::fast_erf_float_expr(arg, bits)`).
+/// WebGPU backends lower `erf`.
pub fn fast_erf_float_expr(arg: impl IntoExpr<PrimExpr>, bits: u8) ->
Result<PrimExpr> {
let ty = PrimType::float(bits);
- // The C++ coefficients are `float` literals.
+ // The coefficients are single-precision values.
let constant =
|value: f32| -> ExprResult { FloatImm::new(&ty, f64::from(value),
UNKNOWN_LOC).into() };
// alpha_1, alpha_3, ..., alpha_13: the odd coefficients of the numerator.
@@ -335,7 +337,7 @@ pub fn fast_erf_float_expr(arg: impl IntoExpr<PrimExpr>,
bits: u8) -> Result<Pri
];
let x = arg.into_owned_expr()?.min(constant(4.0)).max(constant(-4.0));
let x2 = &x * &x;
- // Horner's rule from the highest coefficient down, as C++ writes it out.
+ // Horner's rule from the highest coefficient down.
let mut p = &x2 * constant(alpha[6]) + constant(alpha[5]);
for coefficient in alpha[..5].iter().rev() {
p = &x2 * p + constant(*coefficient);
diff --git a/rust/src/tirx/op_attr.rs b/rust/src/tirx/op_attr.rs
index 1d15b1287..590b2610b 100644
--- a/rust/src/tirx/op_attr.rs
+++ b/rust/src/tirx/op_attr.rs
@@ -17,14 +17,14 @@
* under the License.
*/
-//! The operator attributes of TIRx (C++ `tvm::tirx::op_attr`), for
-//! [`OpDef::set_attr`](crate::ir::OpDef::set_attr).
+//! Operator attributes. The typed keys of the attributes the dialect reads
from an operator, set
+//! with [`OpDef::set_attr`](crate::ir::OpDef::set_attr).
use tvm_ffi::String;
use crate::ir::AttrKey;
-/// The global symbol of the operator after lowering (C++ `kGlobalSymbol`,
`TGlobalSymbol`).
+/// The global symbol the operator lowers to.
pub const GLOBAL_SYMBOL: AttrKey<String> = AttrKey::new("TGlobalSymbol");
-/// Whether the operator is overloaded for vector forms (C++ `kVectorizable`,
`TVectorizable`).
+/// Whether the operator is overloaded for vector operands.
pub const VECTORIZABLE: AttrKey<bool> = AttrKey::new("TVectorizable");
diff --git a/rust/src/tirx/structural.rs b/rust/src/tirx/structural.rs
index 19b5991f8..bfc48d411 100644
--- a/rust/src/tirx/structural.rs
+++ b/rust/src/tirx/structural.rs
@@ -17,12 +17,9 @@
* under the License.
*/
-//! Structural traversal with the facts in scope.
-//!
-//! A [`StructuralVisitorWithConstraintContext`] visits values and a
-//! [`StructuralMutatorWithConstraintContext`] transforms them, keeping the
[`ConstraintContext`] of
-//! the bindings, ranges and guards in effect at the current value. From Rust,
run them with the
-//! `structural_*_with_constraint_context` entry points.
+//! Constraint-aware traversal. A [`StructuralVisitorWithConstraintContext`]
visits values and a
+//! [`StructuralMutatorWithConstraintContext`] rewrites them, each keeping the
[`ConstraintContext`]
+//! of the bindings, ranges and guards in scope at the current value.
// tvm-ffi-stubgen(missing-docs): allow
@@ -53,6 +50,9 @@ struct FfiStructuralVisitorObj {
}
// tvm-ffi-stubgen(end)
+/// A structural mutator that keeps the [`ConstraintContext`] in scope at the
current value, read
+/// with `constraints`; Rust code runs one with
+///
[`structural_mutate_with_constraint_context`](mod@super::structural_mutate_with_constraint_context).
// tvm-ffi-stubgen(begin): object/tirx.StructuralMutatorWithConstraintContext
#[allow(missing_docs)]
#[repr(C)]
@@ -80,6 +80,9 @@ impl StructuralMutatorWithConstraintContextObj {
}
// tvm-ffi-stubgen(end)
+/// A structural visitor that keeps the [`ConstraintContext`] in scope at the
current value, read
+/// with `constraints`; Rust code runs one with
+///
[`structural_visit_with_constraint_context`](mod@super::structural_visit_with_constraint_context).
// tvm-ffi-stubgen(begin): object/tirx.StructuralVisitorWithConstraintContext
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/tirx/structural_map_with_constraint_context.rs
b/rust/src/tirx/structural_map_with_constraint_context.rs
index 854ea836b..26b6872b6 100644
--- a/rust/src/tirx/structural_map_with_constraint_context.rs
+++ b/rust/src/tirx/structural_map_with_constraint_context.rs
@@ -17,7 +17,9 @@
* under the License.
*/
-//! Constraint-aware structural mapping: rewrite each value with the facts in
scope.
+//! Constraint-aware structural mapping.
[`structural_map_with_constraint_context`] calls a closure
+//! on every value with the facts in scope there and replaces the value with
what it returns,
+//! unless that is `Unchanged`.
use std::borrow::BorrowMut;
diff --git a/rust/src/tirx/structural_mutate_with_constraint_context.rs
b/rust/src/tirx/structural_mutate_with_constraint_context.rs
index cfb336c72..b0447359b 100644
--- a/rust/src/tirx/structural_mutate_with_constraint_context.rs
+++ b/rust/src/tirx/structural_mutate_with_constraint_context.rs
@@ -17,7 +17,9 @@
* under the License.
*/
-//! Constraint-aware structural mutation: rewrite values with the facts in
scope.
+//! Constraint-aware structural mutation.
[`structural_mutate_with_constraint_context`] rewrites a
+//! value with a dispatcher whose handlers read the facts in scope from the
+//! [`StructuralMutatorWithConstraintContextView`] they are given.
use std::borrow::BorrowMut;
use std::cell::UnsafeCell;
diff --git a/rust/src/tirx/structural_visit_with_constraint_context.rs
b/rust/src/tirx/structural_visit_with_constraint_context.rs
index a46081a6d..634247087 100644
--- a/rust/src/tirx/structural_visit_with_constraint_context.rs
+++ b/rust/src/tirx/structural_visit_with_constraint_context.rs
@@ -17,7 +17,9 @@
* under the License.
*/
-//! Constraint-aware structural visiting: visit values with the facts in scope.
+//! Constraint-aware structural visiting.
[`structural_visit_with_constraint_context`] visits a
+//! value with a dispatcher whose handlers read the facts in scope from the
+//! [`StructuralVisitorWithConstraintContextView`] they are given.
use std::borrow::BorrowMut;
use std::cell::UnsafeCell;
diff --git a/rust/src/tirx/structural_walk_with_constraint_context.rs
b/rust/src/tirx/structural_walk_with_constraint_context.rs
index c74e1d3cb..f69d07d07 100644
--- a/rust/src/tirx/structural_walk_with_constraint_context.rs
+++ b/rust/src/tirx/structural_walk_with_constraint_context.rs
@@ -17,7 +17,8 @@
* under the License.
*/
-//! Constraint-aware structural walking: observe values with the facts in
scope.
+//! Constraint-aware structural walking.
[`structural_walk_with_constraint_context`] calls a closure
+//! on every value with the facts in scope there, to observe the value, not to
rewrite it.
use std::borrow::BorrowMut;
use tvm_ffi::{
diff --git a/rust/src/tirx/tensor.rs b/rust/src/tirx/tensor.rs
index f8da86bf3..d28bc7516 100644
--- a/rust/src/tirx/tensor.rs
+++ b/rust/src/tirx/tensor.rs
@@ -17,35 +17,13 @@
* under the License.
*/
-//! Tensor types.
-//!
-//! A [`TensorType`] is the element type, memory scope, shape, data alignment
and optional layout of
-//! a tensor.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::{Expr, PrimType, PrimVar, Subscript, Var, UNKNOWN_LOC};
-//! use tvm_tirx::runtime::MemoryScope;
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::tirx::TensorType;
-//!
-//! let scratch = MemoryScope::register("my_ext.scratch")?;
-//! let ty = TensorType::new([16, 16], PrimType::float(32), &scratch, 0, None,
UNKNOWN_LOC)?;
-//! assert_eq!(script(&ty)?, r#"T.Tensor((16, 16), T.float32,
scope="my_ext.scratch")"#);
-//!
-//! let a = Expr::from(Var::new("A", &ty, UNKNOWN_LOC));
-//! let i = PrimVar::new("i", PrimType::int(32), UNKNOWN_LOC);
-//! let point = [Subscript::Index(i.into()), Subscript::Index(0.into())];
-//! let element = a.subscript(point, UNKNOWN_LOC)?;
-//! assert_eq!(script(&Expr::try_from(element)?)?, "A[i, 0]");
-//!
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Tensors. A tensor is a [`Var`] whose type is a [`TensorType`], which fixes
its shape, element
+//! type, memory scope, alignment and [`Layout`]; loads and stores subscript
the variable, and
+//! [`TensorVar`] matches one in an operator signature.
// tvm-ffi-stubgen(missing-docs): allow
-// `custom-new` reserves `new` for a constructor with the C++ one's semantics
(defaults); the
-// generated `from_fields` allocates the fields as given.
+// `custom-new` reserves `new` for a hand-written constructor (defaults); the
generated
+// `from_fields` allocates the fields as given.
// tvm-ffi-stubgen(custom-new): tirx.TensorType
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -87,6 +65,32 @@ use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
// tvm-ffi-stubgen(end)
+/// The type of a tensor of `shape` and element type `dtype` in the memory
`scope`, with its data
+/// pointer aligned to `data_alignment` bytes and an optional `layout` from
indices to storage.
+///
+/// The shape is the one loads and stores index. The data pointer is not part
of the type; the
+/// definition of the tensor variable binds it.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::{Expr, PrimType, PrimVar, Subscript, Var, UNKNOWN_LOC};
+/// use tvm_tirx::runtime::MemoryScope;
+/// use tvm_tirx::script::script;
+/// use tvm_tirx::tirx::TensorType;
+///
+/// let scratch = MemoryScope::register("my_ext.scratch")?;
+/// let ty = TensorType::new([16, 16], PrimType::float(32), &scratch, 0, None,
UNKNOWN_LOC)?;
+/// assert_eq!(script(&ty)?, r#"T.Tensor((16, 16), T.float32,
scope="my_ext.scratch")"#);
+///
+/// let a = Expr::from(Var::new("A", &ty, UNKNOWN_LOC));
+/// let i = PrimVar::new("i", PrimType::int(32), UNKNOWN_LOC);
+/// let point = [Subscript::Index(i.into()), Subscript::Index(0.into())];
+/// let element = a.subscript(point, UNKNOWN_LOC)?;
+/// assert_eq!(script(&Expr::try_from(element)?)?, "A[i, 0]");
+///
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/tirx.TensorType
#[allow(missing_docs)]
#[repr(C)]
@@ -169,8 +173,7 @@ tvm_ffi::unsafe_impl_object_ref_upcast!(TensorType => Type);
impl TensorType {
/// Create the type of a tensor of `shape` and `dtype` in `scope`, aligned
to `data_alignment`
- /// bytes (the runtime's allocation alignment when 0 or less), with
`layout` if given (C++
- /// `TensorType(shape, dtype, scope, data_alignment, layout, loc)`).
+ /// bytes (the runtime's allocation alignment when 0 or less), with
`layout` if given.
pub fn new(
shape: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
dtype: impl Into<PrimType>,
@@ -206,8 +209,8 @@ impl TensorType {
}
}
-/// A signature descriptor ([`sig::ArgType`]) of a `Var`, not a subclass,
whose type is a
-/// [`TensorType`] (C++ `tirx::TensorVar`).
+/// A signature descriptor ([`sig::ArgType`]) that matches a [`Var`] whose
type is a
+/// [`TensorType`]; it is not a type of its own.
pub struct TensorVar;
impl sig::ArgType for TensorVar {
@@ -233,8 +236,8 @@ pub fn tensor_is_scalar(tensor: impl IntoObjectRef<Var>,
alloc_or_decl: bool) ->
.try_into()
}
-/// The runtime's allocation alignment (C++ `runtime::kAllocAlignment`, which
a build may set),
-/// read from the default the reflection registers for
`TensorType::data_alignment`.
+/// The runtime's allocation alignment, which a build may configure, read from
the default the
+/// reflection registers for `TensorType::data_alignment`.
fn alloc_alignment() -> Result<i32> {
static ALLOC_ALIGNMENT: std::sync::OnceLock<i32> =
std::sync::OnceLock::new();
if let Some(alignment) = ALLOC_ALIGNMENT.get() {
diff --git a/rust/src/tirx/thread_scope.rs b/rust/src/tirx/thread_scope.rs
index 1bc84f443..5d0d69fd0 100644
--- a/rust/src/tirx/thread_scope.rs
+++ b/rust/src/tirx/thread_scope.rs
@@ -17,22 +17,8 @@
* under the License.
*/
-//! Thread scopes: the levels of a backend's thread hierarchy.
-//!
-//! A [`ThreadScope`], such as `cuda.warp`, is looked up by name with
[`ThreadScope::get`].
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::tirx::ThreadScope;
-//!
-//! let warp = ThreadScope::get("cuda.warp")?;
-//! assert_eq!((warp.backend()?, warp.local_name()?), ("cuda".into(),
"warp".into()));
-//!
-//! let cta = ThreadScope::get("cuda.cta")?;
-//! assert!(cta.is_higher_than(&warp)?);
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Thread scopes. A [`ThreadScope`] is one level of a backend's thread
hierarchy, such as
+//! `cuda.warp`; each backend registers its levels, ranked from the widest
down.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -60,6 +46,23 @@ struct FfiEnumObj {
}
// tvm-ffi-stubgen(end)
+/// A level of a backend's thread hierarchy, named `backend.level`, such as
`cuda.warp`.
+///
+/// Backends register their scopes; [`ThreadScope::get`] only looks one up.
Within a backend,
+/// [`rank`](ThreadScope::rank) orders the levels, a wider scope ranking lower.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::tirx::ThreadScope;
+///
+/// let warp = ThreadScope::get("cuda.warp")?;
+/// assert_eq!((warp.backend()?, warp.local_name()?), ("cuda".into(),
"warp".into()));
+///
+/// let cta = ThreadScope::get("cuda.cta")?;
+/// assert!(cta.is_higher_than(&warp)?);
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/tirx.ThreadScope
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/tirx/tile.rs b/rust/src/tirx/tile.rs
index 153974159..1144fe2c4 100644
--- a/rust/src/tirx/tile.rs
+++ b/rust/src/tirx/tile.rs
@@ -17,8 +17,8 @@
* under the License.
*/
-//! Tile instructions: their attributes ([`instruction`]) and the context they
are dispatched in
-//! ([`dispatch`]).
+//! Tile instructions. [`instruction`] holds the attributes their calls carry
and their validation,
+//! and [`dispatch`] the context a lowering dispatches each call in.
pub mod dispatch;
pub mod instruction;
diff --git a/rust/src/tirx/tile/dispatch.rs b/rust/src/tirx/tile/dispatch.rs
index 45c3aaceb..ab05b394c 100644
--- a/rust/src/tirx/tile/dispatch.rs
+++ b/rust/src/tirx/tile/dispatch.rs
@@ -17,34 +17,9 @@
* under the License.
*/
-//! The context of tile instruction dispatch.
-//!
-//! A [`DispatchContext`] gives a lowering its target, thread scope, launch
parameters and variable
-//! ranges, and the state shared between dispatches.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::{Evaluate, IntImm, PrimType, PrimVar, UNKNOWN_LOC};
-//! use tvm_tirx::target::Target;
-//! use tvm_tirx::tirx::tile::DispatchContext;
-//! use tvm_tirx::tirx::ThreadScope;
-//! use tvm_tirx::tvm_ffi::Map;
-//!
-//! let tx = PrimVar::new("tx", PrimType::int(32), UNKNOWN_LOC);
-//! let ctx = DispatchContext::new(
-//! Target::new("cuda")?,
-//! ThreadScope::get("cuda.warp")?,
-//! [("threadIdx.x", tx, 32.into())],
-//! &Map::new(), false, &Map::new(), &Map::new(), &Map::new(), &Map::new(),
-//! "warp",
-//! )?;
-//!
-//! ctx.cache_set("my_ext.zero", &IntImm::int32(0))?;
-//! assert!(ctx.cache_get("my_ext.zero")?.is_some());
-//! ctx.add_init_stmt(Evaluate::new(1, UNKNOWN_LOC), false)?;
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Tile instruction dispatch. A lowering dispatches each tile instruction
call with a
+//! [`DispatchContext`], which describes where the call runs and collects the
statements and
+//! allocations the dispatch requests.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -83,6 +58,38 @@ use tvm_ffi::String;
use tvm_ffi::object::ObjectRef;
// tvm-ffi-stubgen(end)
+/// The context a tile instruction call is lowered in: the `target`, the
`thread_scope` of the
+/// call, the kernel's `launch_params`, the ranges of the loop variables in
scope
+/// (`var_range_map`), the axis views across and within teams (`inter`,
`intra`) and the
+/// `scope_kind`.
+///
+/// Through it a dispatch requests allocations and statements, and caches
objects that later
+/// dispatches of the same lowering reuse. A context with `alloc_only` set is
used only to allocate
+/// tensors.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::{Evaluate, IntImm, PrimType, PrimVar, UNKNOWN_LOC};
+/// use tvm_tirx::target::Target;
+/// use tvm_tirx::tirx::tile::DispatchContext;
+/// use tvm_tirx::tirx::ThreadScope;
+/// use tvm_tirx::tvm_ffi::Map;
+///
+/// let tx = PrimVar::new("tx", PrimType::int(32), UNKNOWN_LOC);
+/// let ctx = DispatchContext::new(
+/// Target::new("cuda")?,
+/// ThreadScope::get("cuda.warp")?,
+/// [("threadIdx.x", tx, 32.into())],
+/// &Map::new(), false, &Map::new(), &Map::new(), &Map::new(), &Map::new(),
+/// "warp",
+/// )?;
+///
+/// ctx.cache_set("my_ext.zero", &IntImm::int32(0))?;
+/// assert!(ctx.cache_get("my_ext.zero")?.is_some());
+/// ctx.add_init_stmt(Evaluate::new(1, UNKNOWN_LOC), false)?;
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/tirx.DispatchContext
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/tirx/tile/instruction.rs
b/rust/src/tirx/tile/instruction.rs
index 06ebb3308..ab939455c 100644
--- a/rust/src/tirx/tile/instruction.rs
+++ b/rust/src/tirx/tile/instruction.rs
@@ -19,10 +19,9 @@
// The generated constructors take one argument per reflected field.
#![allow(clippy::too_many_arguments)]
-//! The attributes of tile instruction calls.
-//!
-//! [`ScopeAttrs`] holds the thread scope of an instruction and [`TrnAttrs`]
the scope and options
-//! of a Trainium (`trn`) instruction; their type keys are `tirx.tensor.*`.
+//! Tile instruction calls. Each call carries the static attributes its
operator registers, such as
+//! [`ScopeAttrs`] or [`TrnAttrs`], and is checked by the validator
+//! [`configure_tensor_instruction`] installs.
// tvm-ffi-stubgen(missing-docs): allow
// Every object registered under `tirx.tensor` has its block in this file.
@@ -54,6 +53,7 @@ use tvm_ffi::Result;
use tvm_ffi::String;
// tvm-ffi-stubgen(end)
+/// The attributes of a tile instruction call that needs only the thread
`scope` it runs at.
// tvm-ffi-stubgen(begin): object/tirx.tensor.ScopeAttrs
#[allow(missing_docs)]
#[repr(C)]
@@ -109,6 +109,9 @@ unsafe impl tvm_ffi::object::ImmutableObject for
ScopeAttrsObj {}
tvm_ffi::unsafe_impl_object_ref_upcast!(ScopeAttrs => Attrs);
// tvm-ffi-stubgen(end)
+/// The attributes of a Trainium (`trn`) tile instruction call: the thread
`scope` it runs at, an
+/// optional instruction size limit `max_inst_size`, and the options the `trn`
backend lowers it
+/// with (`opcode`, `reduce_op`, `op0`, `op1`, `axes`, `negate`, `reverse1`).
// tvm-ffi-stubgen(begin): object/tirx.tensor.TrnAttrs
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/tirx/transform.rs b/rust/src/tirx/transform.rs
index 212509c00..5da6d2856 100644
--- a/rust/src/tirx/transform.rs
+++ b/rust/src/tirx/transform.rs
@@ -17,7 +17,8 @@
* under the License.
*/
-//! The TIRx passes ([`mod@transform`]) and function passes
([`function_pass`]).
+//! Passes of the dialect. [`mod@transform`] holds the passes over TIRx
functions, and
+//! [`function_pass`] builds one from a Rust closure.
// tvm-ffi-stubgen(missing-docs): allow
// An object registered under `tirx.transform` that no file has a block for
yet is generated here,
diff --git a/rust/src/tirx/transform/function_pass.rs
b/rust/src/tirx/transform/function_pass.rs
index fae9161eb..839837b53 100644
--- a/rust/src/tirx/transform/function_pass.rs
+++ b/rust/src/tirx/transform/function_pass.rs
@@ -17,34 +17,8 @@
* under the License.
*/
-//! Function passes.
-//!
-//! A [`FunctionPass`] is a [`Pass`] that transforms each TIRx function of a
module.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::transform::{Pass, PassInfo};
-//! use tvm_tirx::tirx::transform::FunctionPass;
-//! # use tvm_tirx::ir::{BaseFunc, DictAttrs, Evaluate, GlobalVar, IRModule,
MissingType, SeqStmt};
-//! # use tvm_tirx::ir::{Var, UNKNOWN_LOC};
-//! # use tvm_tirx::tirx::Function;
-//! # use tvm_tirx::tvm_ffi::Map;
-//! # let body = SeqStmt::new([Evaluate::new(1, UNKNOWN_LOC)], UNKNOWN_LOC);
-//! # let attrs = DictAttrs::new(Map::new());
-//! # let main = Function::new(Vec::<Var>::new(), &body, MissingType::new(),
attrs, UNKNOWN_LOC);
-//! # let main_var = GlobalVar::new("main", UNKNOWN_LOC);
-//! # let functions = Map::from_iter([(main_var, BaseFunc::from(&main))]);
-//! # let module = IRModule::new(&functions, DictAttrs::new(Map::new()),
&Map::new())?;
-//!
-//! let pass = FunctionPass::new(
-//! |func, _module, _ctx| Ok((!func.params.is_empty()).then_some(func)),
-//! PassInfo::new(0, "my_ext.RemoveNullary"),
-//! );
-//! let module = Pass::from(&pass).run(module)?;
-//! assert!(!module.contains("main")?);
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Function passes. A [`FunctionPass`] is a [`Pass`] that rewrites each TIRx
function of a module
+//! on its own, given the module and the pass context.
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -71,6 +45,33 @@ use tvm_ffi::RValueRef;
use tvm_ffi::Result;
// tvm-ffi-stubgen(end)
+/// A pass that replaces each TIRx function of a module by what a function
returns for it, or
+/// removes it for `None`; [`FunctionPass::new`] creates one from a Rust
closure.
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::ir::transform::{Pass, PassInfo};
+/// use tvm_tirx::tirx::transform::FunctionPass;
+/// # use tvm_tirx::ir::{BaseFunc, DictAttrs, Evaluate, GlobalVar, IRModule,
MissingType, SeqStmt};
+/// # use tvm_tirx::ir::{Var, UNKNOWN_LOC};
+/// # use tvm_tirx::tirx::Function;
+/// # use tvm_tirx::tvm_ffi::Map;
+/// # let body = SeqStmt::new([Evaluate::new(1, UNKNOWN_LOC)], UNKNOWN_LOC);
+/// # let attrs = DictAttrs::new(Map::new());
+/// # let main = Function::new(Vec::<Var>::new(), &body, MissingType::new(),
attrs, UNKNOWN_LOC);
+/// # let main_var = GlobalVar::new("main", UNKNOWN_LOC);
+/// # let functions = Map::from_iter([(main_var, BaseFunc::from(&main))]);
+/// # let module = IRModule::new(&functions, DictAttrs::new(Map::new()),
&Map::new())?;
+///
+/// let pass = FunctionPass::new(
+/// |func, _module, _ctx| Ok((!func.params.is_empty()).then_some(func)),
+/// PassInfo::new(0, "my_ext.RemoveNullary"),
+/// );
+/// let module = Pass::from(&pass).run(module)?;
+/// assert!(!module.contains("main")?);
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
// tvm-ffi-stubgen(begin): object/tirx.FunctionPass
#[allow(missing_docs)]
#[repr(C)]
diff --git a/rust/src/tirx/transform/transform.rs
b/rust/src/tirx/transform/transform.rs
index d3a6003b4..da363812a 100644
--- a/rust/src/tirx/transform/transform.rs
+++ b/rust/src/tirx/transform/transform.rs
@@ -17,31 +17,9 @@
* under the License.
*/
-//! The TIRx passes and their configurations.
-//!
-//! Each function returns a [`Pass`]; [`RemoveNoOpConfig`] and
[`UnrollLoopConfig`] configure
+//! The passes of the dialect. Each function returns a [`Pass`] over the TIRx
functions of a module;
+//! [`RemoveNoOpConfig`] and [`UnrollLoopConfig`] are the pass context options
that configure
//! [`remove_no_op`] and [`unroll_loop`].
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::tirx::{transform, Function};
-//! use tvm_tirx::tvm_ffi::{Error, ObjectRefCore};
-//! # use tvm_tirx::ir::{BaseFunc, DictAttrs, Evaluate, GlobalVar, IRModule,
MissingType, SeqStmt};
-//! # use tvm_tirx::ir::{Var, UNKNOWN_LOC};
-//! # use tvm_tirx::tvm_ffi::Map;
-//! # let body = SeqStmt::new([Evaluate::new(1, UNKNOWN_LOC)], UNKNOWN_LOC);
-//! # let attrs = DictAttrs::new(Map::new());
-//! # let main = Function::new(Vec::<Var>::new(), &body, MissingType::new(),
attrs, UNKNOWN_LOC);
-//! # let main_var = GlobalVar::new("main", UNKNOWN_LOC);
-//! # let functions = Map::from_iter([(main_var, BaseFunc::from(&main))]);
-//! # let module = IRModule::new(&functions, DictAttrs::new(Map::new()),
&Map::new())?;
-//!
-//! let module = transform::remove_no_op()?.run(module)?;
-//! let main: Function = module.lookup_name("main")?.try_into()?;
-//! assert!(main.body.seq.is_empty());
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
// tvm-ffi-stubgen(missing-docs): allow
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
@@ -62,6 +40,10 @@ use tvm_ffi::Result;
use tvm_ffi::String;
// tvm-ffi-stubgen(end)
+/// The configuration of [`remove_no_op`], the pass context option
`tirx.RemoveNoOp`:
+/// `ignore_profiler_call` treats profiler calls as no-ops, and a non-zero
+/// `max_simplification_steps` makes the simplifier fail after that many
rewrite steps, for
+/// debugging.
// tvm-ffi-stubgen(begin): object/tirx.transform.RemoveNoOpConfig
#[allow(missing_docs)]
#[repr(C)]
@@ -109,6 +91,12 @@ impl RemoveNoOpConfig {
unsafe impl tvm_ffi::object::ImmutableObject for RemoveNoOpConfigObj {}
// tvm-ffi-stubgen(end)
+/// The configuration of [`unroll_loop`], the pass context option
`tirx.UnrollLoop`.
+///
+/// A loop is unrolled automatically within `auto_max_step` steps,
`auto_max_depth` nesting levels
+/// and an extent of `auto_max_extent`; `explicit_unroll` unrolls the loop in
the IR rather than
+/// marking it with a pragma, and `unroll_local_access` always unrolls loops
that access local
+/// memory.
// tvm-ffi-stubgen(begin): object/tirx.transform.UnrollLoopConfig
#[allow(missing_docs)]
#[repr(C)]
@@ -307,6 +295,27 @@ pub fn pointer_value_type_rewrite() -> Result<Pass> {
}
/// Remove effect-free statements and provably empty control flow
(`tirx.transform.RemoveNoOp`).
+///
+/// # Examples
+///
+/// ```
+/// use tvm_tirx::tirx::{transform, Function};
+/// use tvm_tirx::tvm_ffi::{Error, ObjectRefCore};
+/// # use tvm_tirx::ir::{BaseFunc, DictAttrs, Evaluate, GlobalVar, IRModule,
MissingType, SeqStmt};
+/// # use tvm_tirx::ir::{Var, UNKNOWN_LOC};
+/// # use tvm_tirx::tvm_ffi::Map;
+/// # let body = SeqStmt::new([Evaluate::new(1, UNKNOWN_LOC)], UNKNOWN_LOC);
+/// # let attrs = DictAttrs::new(Map::new());
+/// # let main = Function::new(Vec::<Var>::new(), &body, MissingType::new(),
attrs, UNKNOWN_LOC);
+/// # let main_var = GlobalVar::new("main", UNKNOWN_LOC);
+/// # let functions = Map::from_iter([(main_var, BaseFunc::from(&main))]);
+/// # let module = IRModule::new(&functions, DictAttrs::new(Map::new()),
&Map::new())?;
+///
+/// let module = transform::remove_no_op()?.run(module)?;
+/// let main: Function = module.lookup_name("main")?.try_into()?;
+/// assert!(main.body.seq.is_empty());
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
pub fn remove_no_op() -> Result<Pass> {
tvm_ffi::cached_global_func!("tirx.transform.RemoveNoOp").call_tuple(())?.try_into()
}
diff --git a/rust/src/tirx/type.rs b/rust/src/tirx/type.rs
index 5fa578238..2547cda8c 100644
--- a/rust/src/tirx/type.rs
+++ b/rust/src/tirx/type.rs
@@ -17,22 +17,8 @@
* under the License.
*/
-//! Types of TIRx operations.
-//!
-//! [`TensorMapType`] is the type of a tensor map.
-//!
-//! # Examples
-//!
-//! ```
-//! use tvm_tirx::ir::UNKNOWN_LOC;
-//! use tvm_tirx::script::script;
-//! use tvm_tirx::tirx::TensorMapType;
-//!
-//! let ty = TensorMapType::new(UNKNOWN_LOC);
-//! assert_eq!(script(&ty)?, "T.TensorMap()");
-//!
-//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-//! ```
+//! Types of the dialect beyond tensors. A [`TensorMapType`] types the
tensor-map descriptors that
+//! backend memory operations take; tensors are typed by
[`TensorType`](super::TensorType).
// tvm-ffi-stubgen(missing-docs): allow
@@ -44,6 +30,10 @@ use std::ops::Deref;
use tvm_ffi::ObjectArc;
// tvm-ffi-stubgen(end)
+/// The type of a tensor-map descriptor, which backend memory operations take.
+///
+/// A pointer to it points at the descriptor storage, not at tensor data; the
operation that
+/// creates a descriptor defines its shape and layout.
// tvm-ffi-stubgen(begin): object/tirx.TensorMapType
#[allow(missing_docs)]
#[repr(C)]
