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 b52e19f2b [Rust] Align the remaining constructors,
primitive-expression types and `Result`s with C++ (#212)
b52e19f2b is described below
commit b52e19f2b47e68acbc5821afc2d12eda6d707d0e
Author: Linzhang Li <[email protected]>
AuthorDate: Fri Nov 6 12:33:12 2026 -0500
[Rust] Align the remaining constructors, primitive-expression types and
`Result`s with C++ (#212)
Follow-up to #207 and #205, from an audit of the Rust API before it is
locked down. The changes, in four parts:
1. **Constructors.** #207 made a type's `new` the constructor with the
C++ one's arguments, and named the generated complete-field constructor
`from_fields`. `IRModule` and `PassContext` still had their C++
constructor under another name, and a generated `new` with other
arguments. Their `new` is now the C++ constructor. A type with no C++
constructor keeps its generated `new`.
2. **`PrimExpr` narrowing.** #205 typed primitive expressions the way
C++ does, but it left the layout, index-map and linear-analysis APIs as
`Expr`. They now take and return `PrimExpr` and `PrimVar` wherever C++
does.
3. **`Sequential::new` and the unknown location.** `Sequential::new`
takes the C++ arguments. The unknown location goes from four spellings
to one per use.
4. **Types and `Result` follow C++.** `For.kind` is a `ForKind`,
`Subscript` and `convert` use `PrimExpr`, and three constructors that
cannot fail no longer return `Result`.
## API changes
| Before | After |
| --- | --- |
| `IRModule::from_functions(&functions, None, &infos)` |
`IRModule::new(&functions, DictAttrs::new(Map::new()), &infos)` |
| `IRModule::new(functions, attrs, infos, global_var_map)` |
`IRModule::from_fields(..)` |
| `PassContext::create(opt_level, required, disabled, instruments,
config)` | `PassContext::new(..)`, same arguments |
| `StringType::new(loc)`, `OpaqueType::new(loc)` | `StringType::new()`,
`OpaqueType::new()` |
| `IndexMap::new(Array<Var>, Array<PrimExpr>, None)` |
`IndexMap::new([&i, &j], [&j, &i], None)?` |
| `layout.size(..)` / `layout.apply(..)` returning `Expr` / `Map<String,
Expr>` | `PrimExpr` / `Map<String, PrimExpr>` |
| iterator maps over `Map<Var, Range>` | `Map<PrimVar, Range>` |
| `detect_linear_equation(expr, vars) -> Array<Expr>` | `->
Array<PrimExpr>`, with `vars` as `PrimVar`s |
| `Location::unknown()` | `UnknownLoc::new()` for a value, `UNKNOWN_LOC`
for a `loc` argument |
| `UnknownLoc::from_fields()` | `UnknownLoc::new()` |
| `Sequential::new(passes, opt_level, name)` | `Sequential::new(passes,
PassInfo::new(opt_level, name))` |
| `for_.kind == ForKind::Parallel as i32`, `kind as i32` | `for_.kind ==
ForKind::Parallel`, `kind.as_raw()` |
| `match kind { ForKind::Default => .., ... }` (exhaustive) | add a `_`
arm: `ForKind` is an open newtype |
| `Subscript::Index(expr)` with an `Expr` |
`Subscript::Index(prim_expr)` |
| `convert(..) -> Result<Expr>` | `-> Result<PrimExpr>` |
| `ConstraintContext::new()?`, `SourceName::get(..)?`, `Axis::get(..)?`
| without `?` |
---
rust/src/codegen.rs | 6 +-
rust/src/ir.rs | 2 +-
rust/src/ir/attrs.rs | 2 +-
rust/src/ir/base.rs | 18 +++++-
rust/src/ir/expr.rs | 11 ++--
rust/src/ir/location.rs | 50 ++++------------
rust/src/ir/module.rs | 25 ++++----
rust/src/ir/prim.rs | 6 +-
rust/src/ir/prim/op.rs | 10 +++-
rust/src/ir/stmt.rs | 60 +++++++++++++------
rust/src/ir/transform/core.rs | 34 ++++++-----
rust/src/ir/transform/instrument.rs | 6 +-
rust/src/ir/type.rs | 22 +++++--
rust/src/ir/type_def.rs | 4 +-
rust/src/lib.rs | 4 +-
rust/src/macro/def_expr_op_overloads.rs | 4 +-
rust/src/sym/analyzer.rs | 2 +-
rust/src/sym/constraint.rs | 15 ++---
rust/src/sym/iter_affine_map.rs | 16 ++---
rust/src/tirx/analysis/analysis.rs | 22 +++----
rust/src/tirx/index_map.rs | 43 +++++++++----
rust/src/tirx/layout.rs | 70 +++++++++++++---------
.../structural_mutate_with_constraint_context.rs | 2 +-
.../structural_visit_with_constraint_context.rs | 2 +-
rust/src/tirx/tensor.rs | 4 +-
rust/src/tirx/transform/function_pass.rs | 5 +-
rust/src/tirx/transform/transform.rs | 5 +-
rust/tests/ir/expr.rs | 38 ++++++------
rust/tests/ir/location.rs | 4 +-
rust/tests/ir/main.rs | 6 +-
rust/tests/ir/module.rs | 8 ++-
rust/tests/ir/op_def.rs | 2 +-
rust/tests/ir/prim/expr.rs | 12 ++--
rust/tests/ir/prim/op.rs | 46 +++++++-------
rust/tests/ir/prim/vector_expr.rs | 8 +--
rust/tests/ir/stmt.rs | 24 ++++----
rust/tests/ir/transform/core.rs | 11 ++--
rust/tests/ir/transform/instrument.rs | 4 +-
rust/tests/ir/type.rs | 12 ++--
rust/tests/sym/int_set.rs | 4 +-
rust/tests/sym/iter_affine_map.rs | 2 +-
rust/tests/sym/main.rs | 8 +--
rust/tests/tirx/analysis/analysis.rs | 4 +-
rust/tests/tirx/function.rs | 4 +-
rust/tests/tirx/index_map.rs | 23 ++++---
rust/tests/tirx/layout.rs | 8 +--
rust/tests/tirx/main.rs | 23 ++++---
rust/tests/tirx/op.rs | 6 +-
rust/tests/tirx/tensor.rs | 6 +-
rust/tests/tirx/tile/dispatch.rs | 4 +-
rust/tests/tirx/tile/instruction.rs | 4 +-
51 files changed, 399 insertions(+), 322 deletions(-)
diff --git a/rust/src/codegen.rs b/rust/src/codegen.rs
index b4004ba9c..e8a1aa475 100644
--- a/rust/src/codegen.rs
+++ b/rust/src/codegen.rs
@@ -52,10 +52,10 @@ use tvm_ffi::String;
///
/// ```
/// use tvm_tirx::codegen;
-/// use tvm_tirx::ir::{GlobalVar, IRModule, PrimVar, UNKNOWN_LOC};
+/// use tvm_tirx::ir::{DictAttrs, GlobalVar, IRModule, PrimVar, UNKNOWN_LOC};
/// use tvm_tirx::target::Target;
/// use tvm_tirx::tvm_ffi::Map;
-/// # use tvm_tirx::ir::{DictAttrs, Expr, PrimType, Return, SeqStmt, Var};
+/// # use tvm_tirx::ir::{Expr, PrimType, Return, SeqStmt, Var};
/// # use tvm_tirx::tirx::Function;
/// # use tvm_tirx::tvm_ffi::{Any, String};
/// # let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
@@ -65,7 +65,7 @@ use tvm_ffi::String;
/// # let add_one = Function::new([&x], &body, PrimType::int(32), attrs,
UNKNOWN_LOC);
///
/// let functions = Map::from_iter([(GlobalVar::new("add_one", UNKNOWN_LOC),
add_one.into())]);
-/// let module = IRModule::from_functions(&functions, None, &Map::new())?;
+/// let module = IRModule::new(&functions, DictAttrs::new(Map::new()),
&Map::new())?;
///
/// let compiled = codegen::build_module(&module, Target::new("c")?)?;
/// assert!(compiled.implements_function_with_imports("add_one")?);
diff --git a/rust/src/ir.rs b/rust/src/ir.rs
index 84783051c..50275b275 100644
--- a/rust/src/ir.rs
+++ b/rust/src/ir.rs
@@ -42,7 +42,7 @@
//! let func = Function::new([x], &body, PrimType::int(32), attrs,
UNKNOWN_LOC);
//!
//! let functions = Map::from_iter([(GlobalVar::new("main", UNKNOWN_LOC),
BaseFunc::from(func))]);
-//! let module = IRModule::from_functions(&functions, None, &Map::new())?;
+//! let module = IRModule::new(&functions, DictAttrs::new(Map::new()),
&Map::new())?;
//! assert!(script(&module)?.contains("def main(x: T.int32) -> T.int32:"));
//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
//! ```
diff --git a/rust/src/ir/attrs.rs b/rust/src/ir/attrs.rs
index 006eadefa..49bb77c1b 100644
--- a/rust/src/ir/attrs.rs
+++ b/rust/src/ir/attrs.rs
@@ -33,7 +33,7 @@
//!
//! let dict: Map<String, Any> =
//! [("system_lib_prefix".into(),
Any::from(String::from("my_lib_")))].into_iter().collect();
-//! let module = IRModule::from_functions(&Map::new(),
Some(&DictAttrs::new(dict)), &Map::new())?;
+//! let module = IRModule::new(&Map::new(), DictAttrs::new(dict),
&Map::new())?;
//!
//! let prefix = module.attrs.dict.get(&"system_lib_prefix".into())?.unwrap();
//! assert_eq!(String::try_from(prefix)?, "my_lib_");
diff --git a/rust/src/ir/base.rs b/rust/src/ir/base.rs
index 1b3064012..351a83cdc 100644
--- a/rust/src/ir/base.rs
+++ b/rust/src/ir/base.rs
@@ -31,17 +31,22 @@
//!
//! # let double = Function::from_typed(|x: i64| -> Result<i64> { Ok(x * 2) });
//! # Function::register_global("my_ext.double", double)?;
-//! let double = EnvFunc::new("my_ext.double",
Function::get_global("my_ext.double")?);
+//! let double = EnvFunc::get("my_ext.double")?;
//!
//! assert_eq!(i64::try_from(double.func.call_tuple((21_i64,))?)?, 42);
//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
//! ```
+// Names the hand-written wrappers use: the import section keeps them even
when no generated block
+// does.
+// tvm-ffi-stubgen(import-object): tvm_ffi::Result
+
// tvm-ffi-stubgen(begin): import-section
use std::ops::Deref;
use tvm_ffi::Function;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
+use tvm_ffi::Result;
use tvm_ffi::String;
// tvm-ffi-stubgen(end)
@@ -92,3 +97,14 @@ impl EnvFunc {
unsafe impl tvm_ffi::object::ImmutableObject for EnvFuncObj {}
// tvm-ffi-stubgen(end)
+
+impl EnvFunc {
+ /// Return the global function `name` kept with its name (`ir.EnvFuncGet`).
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if no global function is registered as `name`.
+ pub fn get(name: &str) -> Result<Self> {
+
tvm_ffi::cached_global_func!("ir.EnvFuncGet").call_tuple((String::from(name),))?.try_into()
+ }
+}
diff --git a/rust/src/ir/expr.rs b/rust/src/ir/expr.rs
index b683173c2..668e571e7 100644
--- a/rust/src/ir/expr.rs
+++ b/rust/src/ir/expr.rs
@@ -209,13 +209,14 @@ unsafe impl UnsafeInit for ExprObj {
unsafe impl tvm_ffi::object::ImmutableObject for ExprObj {}
// tvm-ffi-stubgen(end)
-/// An index of [`Expr::subscript`]: an expression, or a `start:stop:step`
slice.
+/// An index of [`Expr::subscript`]: a primitive expression, or a
`start:stop:step` slice (C++
+/// `SubscriptSlice`).
#[derive(Clone)]
pub enum Subscript {
- /// The point index of an expression.
- Index(Expr),
+ /// The point index `index`.
+ Index(PrimExpr),
/// The slice `start:stop:step`, whose missing bounds are `None`.
- Slice(Option<Expr>, Option<Expr>, Option<Expr>),
+ Slice(Option<PrimExpr>, Option<PrimExpr>, Option<PrimExpr>),
}
impl From<&Subscript> for Any {
@@ -1835,7 +1836,7 @@ tvm_ffi::unsafe_impl_object_upcast!(StringImm => Expr,
StringImm => Constant);
impl StringImm {
/// Create the string constant `value`, of type [`StringType`] (C++
`StringImm(value, loc)`).
pub fn new(value: &str, loc: impl Into<Location>) -> Self {
- let ty = StringType::new(UNKNOWN_LOC);
+ let ty = StringType::new();
Self::from_fields(loc.into(), ty, String::from(value))
}
}
diff --git a/rust/src/ir/location.rs b/rust/src/ir/location.rs
index ad29b147c..6292624f2 100644
--- a/rust/src/ir/location.rs
+++ b/rust/src/ir/location.rs
@@ -32,16 +32,16 @@
//! use tvm_tirx::ir::{PrimType, PrimVar, SourceLoc, SourceName, Var};
//! use tvm_tirx::tvm_ffi::ObjectRefCore;
//!
-//! let span = SourceLoc::new(SourceName::get("kernel.py")?, 3, 1, 3, 10);
+//! let span = SourceLoc::new(SourceName::get("kernel.py"), 3, 1, 3, 10);
//!
//! let x = PrimVar::new("x", PrimType::int(32), &span);
//! assert!(x.loc.same_as(&span));
//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
//! ```
-// `custom-new` reserves `new` for a constructor with the native one's
semantics (the shared
-// `UnknownLoc` instance); the generated `from_fields` allocates the fields as
given.
-// tvm-ffi-stubgen(custom-new): ir.UnknownLoc
+// C++ has one `UnknownLoc`, which `UnknownLoc::new` returns: `opaque`
generates no allocator for
+// another one. It has no fields, so it loses no field access.
+// tvm-ffi-stubgen(opaque): ir.UnknownLoc
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
@@ -118,16 +118,8 @@ unsafe impl UnsafeInit for LocationObj {
unsafe impl tvm_ffi::object::ImmutableObject for LocationObj {}
// tvm-ffi-stubgen(end)
-impl Location {
- /// Return the shared [`UnknownLoc`] as a [`Location`] value;
[`UNKNOWN_LOC`] is the same
- /// location as a constructor argument.
- pub fn unknown() -> Self {
- UnknownLoc::new().into()
- }
-}
-
// tvm-ffi-stubgen(begin): object/ir.UnknownLoc
-/// Complete: reflected fields fill [24, 24) exactly.
+/// Opaque: vetoed by directive although the layout is reproducible. Fields
are read through the C ABI getters.
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
#[type_key = "ir.UnknownLoc"]
@@ -136,11 +128,6 @@ pub struct UnknownLocObj {
base: LocationObj,
}
-const _: () = {
- assert!(::core::mem::size_of::<UnknownLocObj>() == 24);
- assert!(::core::mem::align_of::<UnknownLocObj>() == 8);
-};
-
#[repr(C)]
#[derive(tvm_ffi::derive::ObjectRef, Clone)]
pub struct UnknownLoc {
@@ -161,21 +148,6 @@ impl Deref for UnknownLocObj {
}
}
-impl UnknownLocObj {
- pub(crate) fn new() -> Self {
- let base = LocationObj::new();
- Self { base }
- }
-}
-
-impl UnknownLoc {
- /// Lossless complete-field allocation.
- pub fn from_fields() -> Self {
- let obj = UnknownLocObj::new();
- Self { base: ObjectArc::new(obj) }
- }
-}
-
unsafe impl tvm_ffi::object::ImmutableObject for UnknownLocObj {}
tvm_ffi::unsafe_impl_object_upcast!(UnknownLoc => Location);
@@ -197,7 +169,7 @@ impl UnknownLoc {
/// A marker for the unknown location as the `loc` argument of a constructor,
which takes any
/// `impl Into<Location>`: `Var::new("x", ty, UNKNOWN_LOC)`.
///
-/// [`Location::unknown`] is the same location as a value.
+/// [`UnknownLoc::new`] is the same location as a value.
#[derive(Clone, Copy, Debug)]
pub struct Unknown;
@@ -206,7 +178,7 @@ pub const UNKNOWN_LOC: Unknown = Unknown;
impl From<Unknown> for Location {
fn from(_: Unknown) -> Self {
- Location::unknown()
+ UnknownLoc::new().into()
}
}
@@ -394,7 +366,11 @@ unsafe impl tvm_ffi::object::ImmutableObject for
SourceNameObj {}
impl SourceName {
/// Return the source name `name`, interned so that equal names share one
object
/// (`ir.SourceName`).
- pub fn get(name: &str) -> Result<Self> {
-
tvm_ffi::cached_global_func!("ir.SourceName").call_tuple((String::from(name),))?.try_into()
+ pub fn get(name: &str) -> Self {
+ support::infallible(
+ tvm_ffi::cached_global_func!("ir.SourceName")
+ .call_tuple((String::from(name),))
+ .and_then(|any| any.try_into()),
+ )
}
}
diff --git a/rust/src/ir/module.rs b/rust/src/ir/module.rs
index 45b38253a..4c52e23bc 100644
--- a/rust/src/ir/module.rs
+++ b/rust/src/ir/module.rs
@@ -27,9 +27,9 @@
//! # Examples
//!
//! ```
-//! use tvm_tirx::ir::{BaseFunc, GlobalVar, IRModule, UNKNOWN_LOC};
+//! use tvm_tirx::ir::{BaseFunc, DictAttrs, GlobalVar, IRModule, UNKNOWN_LOC};
//! use tvm_tirx::tvm_ffi::{Map, ObjectRefCore};
-//! # use tvm_tirx::ir::{DictAttrs, PrimType, Return, SeqStmt, Var};
+//! # use tvm_tirx::ir::{PrimType, Return, SeqStmt, Var};
//! # use tvm_tirx::tirx::Function;
//! # let body = SeqStmt::new([Return::new(0, UNKNOWN_LOC)], UNKNOWN_LOC);
//! # let attrs = DictAttrs::new(Map::new());
@@ -37,12 +37,15 @@
//!
//! let main = GlobalVar::new("main", UNKNOWN_LOC);
//! let functions = Map::from_iter([(main.clone(), BaseFunc::from(func))]);
-//! let module = IRModule::from_functions(&functions, None, &Map::new())?;
+//! let module = IRModule::new(&functions, DictAttrs::new(Map::new()),
&Map::new())?;
//!
//! assert!(module.lookup(&main)?.same_as(&module.lookup_name("main")?));
//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
//! ```
+// `custom-new` reserves `new` for the C++ constructor, which derives the
global variable map; the
+// generated `from_fields` allocates the fields as given.
+// tvm-ffi-stubgen(custom-new): ir.IRModule
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::Any
@@ -121,7 +124,7 @@ impl IRModuleObj {
impl IRModule {
/// Lossless complete-field allocation.
- pub fn new(
+ pub fn from_fields(
functions: impl Into<Map<GlobalVar, BaseFunc>>,
attrs: impl Into<DictAttrs>,
global_infos: impl Into<Map<String, Array<GlobalInfo>>>,
@@ -139,21 +142,17 @@ impl IRModule {
// tvm-ffi-stubgen(end)
impl IRModule {
- /// Create the module of `functions`, with the attributes `attrs` (empty
if `None`) and the
- /// module-level information `global_infos` (C++ `IRModule(functions,
attrs, global_infos)`).
+ /// Create the module of `functions`, with the attributes `attrs` and the
module-level
+ /// information `global_infos` (C++ `IRModule(functions, attrs,
global_infos)`).
///
/// # Errors
///
/// Returns an `InternalError` if two global variables of `functions` have
the same name.
- pub fn from_functions(
+ pub fn new(
functions: &Map<GlobalVar, BaseFunc>,
- attrs: Option<&DictAttrs>,
+ attrs: impl Into<DictAttrs>,
global_infos: &Map<String, Array<GlobalInfo>>,
) -> Result<Self> {
- let attrs = match attrs {
- Some(attrs) => attrs.clone(),
- None => DictAttrs::new(Map::new()),
- };
let mut names = std::collections::HashSet::new();
let mut global_var_map = Vec::with_capacity(functions.len());
for (var, _) in functions.iter() {
@@ -166,7 +165,7 @@ impl IRModule {
global_var_map.push((name, var));
}
let global_var_map = global_var_map.into_iter().collect::<Map<String,
GlobalVar>>();
- Ok(Self::new(functions.clone(), attrs, global_infos.clone(),
global_var_map))
+ Ok(Self::from_fields(functions.clone(), attrs, global_infos.clone(),
global_var_map))
}
/// Create the module of `functions` and `expr`, a function bound to its
`global_symbol`, or to
diff --git a/rust/src/ir/prim.rs b/rust/src/ir/prim.rs
index 1fdc20b17..9d8795279 100644
--- a/rust/src/ir/prim.rs
+++ b/rust/src/ir/prim.rs
@@ -29,7 +29,7 @@
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
-// tvm-ffi-stubgen(import-object): crate::ir::expr::Expr
+// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimExpr
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
mod const_fold;
@@ -44,15 +44,15 @@ pub use vector_expr::*;
// tvm-ffi-stubgen(import-object): std::borrow::Borrow
// tvm-ffi-stubgen(begin): import-section
-use crate::ir::expr::Expr;
use crate::ir::expr::IntoExpr;
+use crate::ir::expr::PrimExpr;
use std::borrow::Borrow;
use tvm_ffi::Result;
// tvm-ffi-stubgen(end)
/// Return whether `lhs` and `rhs` are the same expression tree, without
remapping variables: use
/// structural equality for that (`prim.expr_deep_equal`).
-pub fn expr_deep_equal(lhs: impl IntoExpr, rhs: impl IntoExpr) -> Result<bool>
{
+pub fn expr_deep_equal(lhs: impl IntoExpr<PrimExpr>, rhs: impl
IntoExpr<PrimExpr>) -> Result<bool> {
tvm_ffi::cached_global_func!("prim.expr_deep_equal")
.call_tuple((lhs.into_expr()?.borrow(), rhs.into_expr()?.borrow()))?
.try_into()
diff --git a/rust/src/ir/prim/op.rs b/rust/src/ir/prim/op.rs
index f0431cabd..8929859bd 100644
--- a/rust/src/ir/prim/op.rs
+++ b/rust/src/ir/prim/op.rs
@@ -884,9 +884,13 @@ pub fn make_const_at(
Ok(Broadcast::new(scalar, lanes, loc)?.into())
}
-/// Convert `value` into an expression: a Rust `bool`, integer or float
becomes a constant of the
-/// default type, and an expression is returned as is (`prim.convert`).
-pub fn convert(value: impl Into<Any>) -> Result<Expr> {
+/// Convert `value` into a primitive expression: a Rust `bool`, integer or
float becomes a constant
+/// of the default type, and a primitive expression is returned as is
(`prim.convert`).
+///
+/// # Errors
+///
+/// Returns a `TypeError` if `value` is not a primitive expression or a
literal.
+pub fn convert(value: impl Into<Any>) -> Result<PrimExpr> {
tvm_ffi::cached_global_func!("prim.convert").call_tuple((value.into(),))?.try_into()
}
diff --git a/rust/src/ir/stmt.rs b/rust/src/ir/stmt.rs
index 72515fdc0..60ab7149e 100644
--- a/rust/src/ir/stmt.rs
+++ b/rust/src/ir/stmt.rs
@@ -63,6 +63,9 @@
// tvm-ffi-stubgen(custom-new): ir.Evaluate
// tvm-ffi-stubgen(custom-new): ir.Return
// tvm-ffi-stubgen(custom-new): ir.While
+// A loop's kind is C++ `ForKind`. The `enum` directive declares it as an open
newtype of the
+// native `int` field, so a value Rust does not name reads back intact.
+// tvm-ffi-stubgen(enum): ir.For.kind -> ForKind(i32) { Default=0, Parallel=1,
Vectorized=2, Unrolled=3 }
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::AnyView
@@ -121,6 +124,7 @@ use tvm_ffi::AnyView;
use tvm_ffi::Array;
use tvm_ffi::BigInt;
use tvm_ffi::DLDataType;
+use tvm_ffi::Error;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
@@ -129,6 +133,7 @@ use tvm_ffi::Result;
use tvm_ffi::ShallowCopy;
use tvm_ffi::String;
use tvm_ffi::UnsafeInit;
+use tvm_ffi::VALUE_ERROR;
use tvm_ffi::reflection::drop_if_defined;
use tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
// tvm-ffi-stubgen(end)
@@ -380,6 +385,33 @@ impl AssertStmt {
}
// tvm-ffi-stubgen(begin): object/ir.For
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[repr(transparent)]
+pub struct ForKind(i32);
+
+#[allow(non_upper_case_globals)]
+impl ForKind {
+ pub const Default: Self = Self(0);
+ pub const Parallel: Self = Self(1);
+ pub const Vectorized: Self = Self(2);
+ pub const Unrolled: Self = Self(3);
+ pub const fn from_raw(value: i32) -> Self {
+ Self(value)
+ }
+ pub const fn as_raw(self) -> i32 {
+ self.0
+ }
+}
+
+impl TryFrom<i64> for ForKind {
+ type Error = Error;
+ fn try_from(value: i64) -> Result<Self> {
+ i32::try_from(value).map(Self).map_err(|_| {
+ Error::new(VALUE_ERROR, &format!("ForKind value {value} does not
fit i32"), "")
+ })
+ }
+}
+
/// Complete: reflected fields fill [32, 88) exactly (alignment padding [68,
72)).
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
@@ -391,7 +423,7 @@ pub struct ForObj {
pub begin: PrimExpr,
pub end: PrimExpr,
pub step: PrimExpr,
- pub kind: i32,
+ pub kind: ForKind,
pub body: SeqStmt,
pub annotations: Map<String, Any>,
}
@@ -428,7 +460,7 @@ impl ForObj {
begin: PrimExpr,
end: PrimExpr,
step: PrimExpr,
- kind: i32,
+ kind: ForKind,
body: SeqStmt,
annotations: Map<String, Any>,
) -> Self {
@@ -445,7 +477,7 @@ impl For {
begin: impl Into<PrimExpr>,
end: impl Into<PrimExpr>,
step: impl Into<PrimExpr>,
- kind: i32,
+ kind: impl Into<ForKind>,
body: impl Into<SeqStmt>,
annotations: impl Into<Map<String, Any>>,
) -> Self {
@@ -455,7 +487,7 @@ impl For {
begin.into(),
end.into(),
step.into(),
- kind,
+ kind.into(),
body.into(),
annotations.into(),
);
@@ -468,19 +500,11 @@ unsafe impl tvm_ffi::object::ImmutableObject for ForObj {}
tvm_ffi::unsafe_impl_object_upcast!(For => Stmt);
// tvm-ffi-stubgen(end)
-/// How a [`For`] loop runs its iterations, as C++ `ir::ForKind`.
-#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
-#[repr(i32)]
-pub enum ForKind {
- /// Sequential iterations.
- #[default]
- Default = 0,
- /// Iterations may run in parallel.
- Parallel = 1,
- /// The body is vectorized.
- Vectorized = 2,
- /// The loop is unrolled.
- Unrolled = 3,
+impl Default for ForKind {
+ /// The ordinary loop, of sequential iterations (C++ `ForKind::kDefault`).
+ fn default() -> Self {
+ Self::Default
+ }
}
impl For {
@@ -572,7 +596,7 @@ impl For {
same_type(&end, "var.ty() == end.ty()")?;
let step = promote(step)?;
same_type(&step, "var.ty() == step.ty()")?;
- Ok(Self::from_fields(loc.into(), var, begin, end, step, kind as i32,
body, annotations))
+ Ok(Self::from_fields(loc.into(), var, begin, end, step, kind, body,
annotations))
}
}
diff --git a/rust/src/ir/transform/core.rs b/rust/src/ir/transform/core.rs
index f914d94da..62cd4771f 100644
--- a/rust/src/ir/transform/core.rs
+++ b/rust/src/ir/transform/core.rs
@@ -29,7 +29,7 @@
//! ```
//! use tvm_tirx::ir::transform::{ModulePass, Pass, PassContext, PassInfo,
PassInstrument};
//! use tvm_tirx::ir::transform::Sequential;
-//! use tvm_tirx::ir::{IRModule, IntImm};
+//! use tvm_tirx::ir::{DictAttrs, IRModule, IntImm};
//! use tvm_tirx::tvm_ffi::Map;
//!
//! let mark = |name: &'static str| {
@@ -37,14 +37,19 @@
//! ModulePass::new(pass, PassInfo::new(0, name))
//! };
//!
-//! let pipeline = Pass::from(Sequential::new([mark("a"), mark("b")], 0,
"pipeline"));
-//! let ctx = PassContext::create(2, [""; 0], ["b"],
Vec::<PassInstrument>::new(), None)?;
+//! let info = PassInfo::new(0, "pipeline");
+//! let pipeline = Pass::from(Sequential::new([mark("a"), mark("b")], info));
+//! let ctx = PassContext::new(2, [""; 0], ["b"],
Vec::<PassInstrument>::new(), None)?;
//! let _scope = ctx.enter()?;
-//! let module = pipeline.run(IRModule::from_functions(&Map::new(), None,
&Map::new())?)?;
+//! let empty = IRModule::new(&Map::new(), DictAttrs::new(Map::new()),
&Map::new())?;
+//! let module = pipeline.run(empty)?;
//! assert!(module.get_attr("a")?.is_some() &&
module.get_attr("b")?.is_none());
//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
//! ```
+// `custom-new` reserves `new` for the checked `transform.PassContext`; the
generated `from_fields`
+// allocates the fields as given.
+// tvm-ffi-stubgen(custom-new): transform.PassContext
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): crate::ir::module::IRModule
@@ -266,20 +271,21 @@ tvm_ffi::unsafe_impl_object_upcast!(Sequential => Pass);
// tvm-ffi-stubgen(end)
impl Sequential {
- /// Create a pass named `name`, of level `opt_level`, that runs `passes`
in order
- /// (`transform.Sequential`).
+ /// Create a pass with the information `pass_info` that runs `passes` in
order
+ /// (C++ `Sequential(passes, pass_info)`).
///
/// It skips the passes the current pass context disables, and those whose
level exceeds the
- /// context's `opt_level` unless the context requires them.
+ /// context's `opt_level` unless the context requires them. C++
`Sequential(passes, name)` is
+ /// `Sequential::new(passes, PassInfo::new(0, name))`.
pub fn new(
passes: impl IntoIterator<Item = impl Into<Pass>>,
- opt_level: i32,
- name: &str,
+ pass_info: impl Into<PassInfo>,
) -> Self {
let passes =
passes.into_iter().map(Into::into).collect::<Array<Pass>>();
+ let pass_info = pass_info.into();
support::infallible(
tvm_ffi::cached_global_func!("transform.Sequential")
- .call_tuple((&passes, opt_level, String::from(name)))
+ .call_tuple((&passes, pass_info.opt_level, &pass_info.name))
.and_then(|any| any.try_into()),
)
}
@@ -333,7 +339,7 @@ impl PassContextObj {
impl PassContext {
/// Lossless complete-field allocation.
- pub fn new(
+ pub fn from_fields(
opt_level: i32,
required_pass: impl Into<Array<String>>,
disabled_pass: impl Into<Array<String>>,
@@ -362,7 +368,7 @@ impl PassContext {
///
/// Returns an `AttributeError` if `config` has an option that is not
registered or a `None`
/// value, or the error of converting a value to its option type.
- pub fn create(
+ pub fn new(
opt_level: i32,
required_pass: impl IntoIterator<Item = impl AsRef<str>>,
disabled_pass: impl IntoIterator<Item = impl AsRef<str>>,
@@ -433,8 +439,8 @@ impl PassContext {
/// ```
/// use tvm_tirx::ir::transform::{PassContext, PassInstrument};
///
-/// let outer = PassContext::create(3, [""; 0], [""; 0],
Vec::<PassInstrument>::new(), None)?;
-/// let inner = PassContext::create(1, [""; 0], [""; 0],
Vec::<PassInstrument>::new(), None)?;
+/// let outer = PassContext::new(3, [""; 0], [""; 0],
Vec::<PassInstrument>::new(), None)?;
+/// let inner = PassContext::new(1, [""; 0], [""; 0],
Vec::<PassInstrument>::new(), None)?;
/// let outer_scope = outer.enter()?;
/// let inner_scope = inner.enter()?;
/// assert_eq!(PassContext::current()?.opt_level, 1);
diff --git a/rust/src/ir/transform/instrument.rs
b/rust/src/ir/transform/instrument.rs
index 355aea5fa..fd8852902 100644
--- a/rust/src/ir/transform/instrument.rs
+++ b/rust/src/ir/transform/instrument.rs
@@ -29,7 +29,7 @@
//! use std::sync::{Arc, Mutex};
//! use tvm_tirx::ir::transform::{BasePassInstrument, ModulePass, Pass,
PassContext, PassInfo};
//! use tvm_tirx::ir::transform::PassInstrumentHooks;
-//! use tvm_tirx::ir::IRModule;
+//! use tvm_tirx::ir::{DictAttrs, IRModule};
//! use tvm_tirx::tvm_ffi::Map;
//!
//! let log = Arc::new(Mutex::new(Vec::new()));
@@ -42,9 +42,9 @@
//! };
//! let tracer = BasePassInstrument::new("tracer", hooks);
//!
-//! let _scope = PassContext::create(2, [""; 0], [""; 0], [tracer],
None)?.enter()?;
+//! let _scope = PassContext::new(2, [""; 0], [""; 0], [tracer],
None)?.enter()?;
//! let noop = Pass::from(ModulePass::new(|module, _ctx| Ok(module),
PassInfo::new(0, "noop")));
-//! noop.run(IRModule::from_functions(&Map::new(), None, &Map::new())?)?;
+//! noop.run(IRModule::new(&Map::new(), DictAttrs::new(Map::new()),
&Map::new())?)?;
//! assert_eq!(*log.lock().unwrap(), ["noop"]);
//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
//! ```
diff --git a/rust/src/ir/type.rs b/rust/src/ir/type.rs
index 95ce36ca7..0dee57515 100644
--- a/rust/src/ir/type.rs
+++ b/rust/src/ir/type.rs
@@ -49,6 +49,8 @@
// tvm-ffi-stubgen(custom-new): ir.TupleType
// tvm-ffi-stubgen(custom-new): ir.FuncType
// tvm-ffi-stubgen(custom-new): ir.PrimType
+// tvm-ffi-stubgen(custom-new): ir.OpaqueType
+// tvm-ffi-stubgen(custom-new): ir.StringType
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::ObjectRefCore
@@ -451,7 +453,7 @@ impl OpaqueTypeObj {
impl OpaqueType {
/// Lossless complete-field allocation.
- pub fn new(loc: impl Into<Location>) -> Self {
+ pub fn from_fields(loc: impl Into<Location>) -> Self {
let obj = OpaqueTypeObj::new(loc.into());
Self { base: ObjectArc::new(obj) }
}
@@ -462,8 +464,12 @@ unsafe impl tvm_ffi::object::ImmutableObject for
OpaqueTypeObj {}
tvm_ffi::unsafe_impl_object_upcast!(OpaqueType => Type);
// tvm-ffi-stubgen(end)
-// `ir.OpaqueType` allocates a fresh type at the unknown location: the
generated
-// `OpaqueType::new(Location::unknown())` is the same.
+impl OpaqueType {
+ /// Create an opaque type at the unknown location (C++ `OpaqueType()`).
+ pub fn new() -> Self {
+ Self::from_fields(UNKNOWN_LOC)
+ }
+}
// tvm-ffi-stubgen(begin): object/ir.PrimType
/// Complete: reflected fields fill [32, 40) exactly (alignment padding [36,
40)).
@@ -652,7 +658,7 @@ impl StringTypeObj {
impl StringType {
/// Lossless complete-field allocation.
- pub fn new(loc: impl Into<Location>) -> Self {
+ pub fn from_fields(loc: impl Into<Location>) -> Self {
let obj = StringTypeObj::new(loc.into());
Self { base: ObjectArc::new(obj) }
}
@@ -663,8 +669,12 @@ unsafe impl tvm_ffi::object::ImmutableObject for
StringTypeObj {}
tvm_ffi::unsafe_impl_object_upcast!(StringType => Type);
// tvm-ffi-stubgen(end)
-// `ir.StringType` allocates a fresh type at the unknown location: the
generated
-// `StringType::new(Location::unknown())` is the same.
+impl StringType {
+ /// Create the string type at the unknown location (C++ `StringType()`).
+ pub fn new() -> Self {
+ Self::from_fields(UNKNOWN_LOC)
+ }
+}
// tvm-ffi-stubgen(begin): object/ir.PtrType
/// Complete: reflected fields fill [32, 48) exactly.
diff --git a/rust/src/ir/type_def.rs b/rust/src/ir/type_def.rs
index 341be901a..2266c7274 100644
--- a/rust/src/ir/type_def.rs
+++ b/rust/src/ir/type_def.rs
@@ -38,7 +38,7 @@ use super::support;
/// # Examples
///
/// ```
-/// use tvm_tirx::ir::{register_type, Location, PrimType, Type, TypeObj, Var,
UNKNOWN_LOC};
+/// use tvm_tirx::ir::{register_type, PrimType, Type, TypeObj, Var,
UNKNOWN_LOC};
/// use tvm_tirx::script::script;
/// use tvm_tirx::tvm_ffi::{self, def_object_ref, rs_object};
///
@@ -53,7 +53,7 @@ use super::support;
///
/// impl PairType {
/// pub fn new(first: impl Into<Type>, second: impl Into<Type>) -> Self {
-/// PairTypeObj::new(TypeObj::new(Location::unknown()), first,
second).into()
+/// PairTypeObj::new(TypeObj::new(UNKNOWN_LOC.into()), first,
second).into()
/// }
/// }
///
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 43d8c8a55..f76be4360 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -100,7 +100,7 @@
//! scope.
//!
//! ```
-//! use tvm_tirx::ir::{Location, PrimType, Type, TypeObj};
+//! use tvm_tirx::ir::{PrimType, Type, TypeObj, UNKNOWN_LOC};
//! use tvm_tirx::tvm_ffi::{self, def_object_ref, rs_object, ObjectRefCore};
//!
//! #[rs_object("my_ext.PairType", final)]
@@ -116,7 +116,7 @@
//!
//! impl PairType {
//! pub fn new(first: impl Into<Type>, second: impl Into<Type>) -> Self {
-//! PairTypeObj::new(TypeObj::new(Location::unknown()), first, second,
"").into()
+//! PairTypeObj::new(TypeObj::new(UNKNOWN_LOC.into()), first, second,
"").into()
//! }
//! }
//!
diff --git a/rust/src/macro/def_expr_op_overloads.rs
b/rust/src/macro/def_expr_op_overloads.rs
index e2332b1c3..9e9ceb36f 100644
--- a/rust/src/macro/def_expr_op_overloads.rs
+++ b/rust/src/macro/def_expr_op_overloads.rs
@@ -30,7 +30,7 @@
/// # Examples
///
/// ```
-/// use tvm_tirx::ir::{Expr, ExprObj, Location, PrimExpr, PrimType};
+/// use tvm_tirx::ir::{Expr, ExprObj, PrimExpr, PrimType, UNKNOWN_LOC};
/// use tvm_tirx::sym::Analyzer;
/// use tvm_tirx::tvm_ffi::{self, def_object_ref, rs_object, ObjectRefCore};
///
@@ -44,7 +44,7 @@
/// impl ThreadIndex {
/// /// The `int32` index of the thread along `axis`.
/// pub fn new(axis: i64) -> Self {
-/// let base = ExprObj::new(Location::unknown(),
PrimType::int(32).into());
+/// let base = ExprObj::new(UNKNOWN_LOC.into(),
PrimType::int(32).into());
/// Self::from(ThreadIndexObj::new(base, axis))
/// }
/// }
diff --git a/rust/src/sym/analyzer.rs b/rust/src/sym/analyzer.rs
index dd11e007c..088fd6ebb 100644
--- a/rust/src/sym/analyzer.rs
+++ b/rust/src/sym/analyzer.rs
@@ -32,7 +32,7 @@
//! use tvm_tirx::sym::{Analyzer, ConstraintContext, RangeConstraint};
//!
//! let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
-//! let context = ConstraintContext::new()?;
+//! 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();
diff --git a/rust/src/sym/constraint.rs b/rust/src/sym/constraint.rs
index 630719bcb..02baba18a 100644
--- a/rust/src/sym/constraint.rs
+++ b/rust/src/sym/constraint.rs
@@ -34,7 +34,7 @@
//! use tvm_tirx::sym::{Analyzer, ConstraintContext, PredicateConstraint,
ProofStrength};
//!
//! let x = PrimVar::new("x", PrimType::int(32), UNKNOWN_LOC);
-//! let outer = ConstraintContext::new()?;
+//! 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())?;
@@ -332,8 +332,9 @@ impl Deref for ConstraintContext {
}
impl ConstraintContext {
- /// Create an empty root context, with the complete native-compatible
layout.
- pub fn new() -> Result<Self> {
+ /// Create an empty root context, with the complete native-compatible
layout (C++
+ /// `ConstraintContext()`).
+ pub fn new() -> Self {
Self::allocate(None, 0)
}
@@ -353,11 +354,11 @@ impl ConstraintContext {
"",
));
}
- Self::allocate(Some(parent.base), parent_prefix_size)
+ Ok(Self::allocate(Some(parent.base), parent_prefix_size))
}
- fn allocate(parent: Option<ObjectArc<ConstraintContextObj>>, prefix: i64)
-> Result<Self> {
- Ok(Self {
+ fn allocate(parent: Option<ObjectArc<ConstraintContextObj>>, prefix: i64)
-> Self {
+ Self {
base: ObjectArc::new(ConstraintContextObj {
base: Object::new(),
parent,
@@ -366,7 +367,7 @@ impl ConstraintContext {
// never install a nonempty Rust allocation in this mutable
field.
constraints: UnsafeCell::new(Array::default()),
}),
- })
+ }
}
}
diff --git a/rust/src/sym/iter_affine_map.rs b/rust/src/sym/iter_affine_map.rs
index b85703838..ee7a32391 100644
--- a/rust/src/sym/iter_affine_map.rs
+++ b/rust/src/sym/iter_affine_map.rs
@@ -34,7 +34,7 @@
//! 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([(Var::from(&x), four.clone()), (Var::from(&y),
four)]);
+//! 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)?;
@@ -57,6 +57,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::support
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimExpr
+// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimVar
// The iterator forms are primitive expressions, as C++ `IterMapExpr` derives
from `PrimExpr`, and
// their marks and padding predicate are `PrimExpr`s.
// tvm-ffi-stubgen(upcast): sym.IterMapExpr -> PrimExpr
@@ -76,6 +77,7 @@ use crate::ir::expr::Expr;
use crate::ir::expr::ExprObj;
use crate::ir::expr::IntoExpr;
use crate::ir::expr::PrimExpr;
+use crate::ir::expr::PrimVar;
use crate::ir::expr::Range;
use crate::ir::expr::Var;
use crate::ir::support;
@@ -381,7 +383,7 @@ pub enum IterMapLevel {
/// When they are not, the `indices` of the result are empty and its `errors`
say why.
pub fn detect_iter_map(
indices: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
- input_iters: &Map<Var, Range>,
+ input_iters: &Map<PrimVar, Range>,
predicate: impl IntoExpr<PrimExpr>,
check_level: IterMapLevel,
analyzer: Option<&Analyzer>,
@@ -406,7 +408,7 @@ pub fn detect_iter_map(
/// its `base` (`sym.NormalizeToIterSum`).
pub fn normalize_to_iter_sum(
index: impl IntoExpr<PrimExpr>,
- input_iters: &Map<Var, Range>,
+ input_iters: &Map<PrimVar, Range>,
analyzer: Option<&Analyzer>,
) -> Result<IterSumExpr> {
tvm_ffi::cached_global_func!("sym.NormalizeToIterSum")
@@ -418,7 +420,7 @@ pub fn normalize_to_iter_sum(
/// iterator maps (`sym.IterMapSimplify`).
pub fn iter_map_simplify(
indices: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
- input_iters: &Map<Var, Range>,
+ input_iters: &Map<PrimVar, Range>,
predicate: impl IntoExpr<PrimExpr>,
check_level: IterMapLevel,
analyzer: Option<&Analyzer>,
@@ -453,8 +455,8 @@ pub fn normalize_iter_map_to_expr(expr: impl
IntoExpr<PrimExpr>) -> Result<PrimE
/// the outer and inner predicates, or empty when the bindings cannot be
divided.
pub fn subspace_divide(
bindings: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
- input_iters: &Map<Var, Range>,
- sub_iters: impl IntoIterator<Item = impl Into<Var>>,
+ input_iters: &Map<PrimVar, Range>,
+ sub_iters: impl IntoIterator<Item = impl Into<PrimVar>>,
predicate: impl IntoExpr<PrimExpr>,
check_level: IterMapLevel,
analyzer: Option<&Analyzer>,
@@ -462,7 +464,7 @@ pub fn subspace_divide(
) -> Result<Array<Array<IterMark>>> {
let bindings =
bindings.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<PrimExpr>>>()?;
- let sub_iters =
sub_iters.into_iter().map(Into::into).collect::<Array<Var>>();
+ let sub_iters =
sub_iters.into_iter().map(Into::into).collect::<Array<PrimVar>>();
let (predicate, analyzer) = (predicate.into_owned_expr()?,
analyzer.cloned());
tvm_ffi::cached_global_func!("sym.SubspaceDivide")
.call_packed(&[
diff --git a/rust/src/tirx/analysis/analysis.rs
b/rust/src/tirx/analysis/analysis.rs
index 31ddb71d5..9c9f12bf9 100644
--- a/rust/src/tirx/analysis/analysis.rs
+++ b/rust/src/tirx/analysis/analysis.rs
@@ -38,7 +38,8 @@
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
-// tvm-ffi-stubgen(import-object): crate::ir::expr::Expr
+// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimExpr
+// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimVar
// tvm-ffi-stubgen(import-object): crate::ir::expr::Var
// tvm-ffi-stubgen(import-object): crate::ir::module::IRModule
// tvm-ffi-stubgen(import-object): crate::tirx::function::Function
@@ -50,8 +51,9 @@
// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
// tvm-ffi-stubgen(begin): import-section
-use crate::ir::expr::Expr;
use crate::ir::expr::IntoExpr;
+use crate::ir::expr::PrimExpr;
+use crate::ir::expr::PrimVar;
use crate::ir::expr::Var;
use crate::ir::module::IRModule;
use crate::tirx::function::Function;
@@ -107,10 +109,10 @@ pub fn check_well_formed_module(
/// Return the coefficients `c` of `expr = sum(vars[i] * c[i]) + c[n]`, each
invariant of `vars`, or
/// an empty array if `expr` has no such form
(`tirx.analysis.DetectLinearEquation`).
pub fn detect_linear_equation(
- expr: impl IntoExpr,
- vars: impl IntoIterator<Item = impl Into<Var>>,
-) -> Result<Array<Expr>> {
- let vars = vars.into_iter().map(Into::into).collect::<Array<Var>>();
+ expr: impl IntoExpr<PrimExpr>,
+ vars: impl IntoIterator<Item = impl Into<PrimVar>>,
+) -> Result<Array<PrimExpr>> {
+ let vars = vars.into_iter().map(Into::into).collect::<Array<PrimVar>>();
tvm_ffi::cached_global_func!("tirx.analysis.DetectLinearEquation")
.call_tuple((expr.into_expr()?.borrow(), &vars))?
.try_into()
@@ -120,10 +122,10 @@ pub fn detect_linear_equation(
/// with `None` for no bound, or an empty array if `expr` is not a clip bound
/// (`tirx.analysis.DetectClipBound`).
pub fn detect_clip_bound(
- expr: impl IntoExpr,
- vars: impl IntoIterator<Item = impl Into<Var>>,
-) -> Result<Array<Option<Expr>>> {
- let vars = vars.into_iter().map(Into::into).collect::<Array<Var>>();
+ expr: impl IntoExpr<PrimExpr>,
+ vars: impl IntoIterator<Item = impl Into<PrimVar>>,
+) -> Result<Array<Option<PrimExpr>>> {
+ let vars = vars.into_iter().map(Into::into).collect::<Array<PrimVar>>();
tvm_ffi::cached_global_func!("tirx.analysis.DetectClipBound")
.call_tuple((expr.into_expr()?.borrow(), &vars))?
.try_into()
diff --git a/rust/src/tirx/index_map.rs b/rust/src/tirx/index_map.rs
index 13c49bc85..6dc704434 100644
--- a/rust/src/tirx/index_map.rs
+++ b/rust/src/tirx/index_map.rs
@@ -25,16 +25,14 @@
//! # Examples
//!
//! ```
-//! use tvm_tirx::ir::{PrimExpr, PrimType, PrimVar, Var, UNKNOWN_LOC};
+//! 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 final_indices = Array::from_vec(vec![PrimExpr::from(&j),
PrimExpr::from(&i)]);
-//! let initial_indices = Array::from_vec(vec![Var::from(i), Var::from(j)]);
-//! let transpose = IndexMap::new(initial_indices, final_indices, None);
+//! let transpose = IndexMap::new([&i, &j], [&j, &i], None)?;
//! assert_eq!(script(&transpose)?, "T.index_map(lambda i, j: (j, i))");
//!
//! let text =
@@ -53,13 +51,20 @@
// 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
+// 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.
// tvm-ffi-stubgen(field): tirx.IndexMap.final_indices -> Array<PrimExpr>
+// Its indices are `PrimVar`s, as C++ `IndexMapObj` holds them.
+// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimVar
+// tvm-ffi-stubgen(import-object): tvm_ffi::Array
+// tvm-ffi-stubgen(field): tirx.IndexMap.initial_indices -> Array<PrimVar>
// tvm-ffi-stubgen(begin): import-section
use crate::ir::expr::IntoExpr;
use crate::ir::expr::PrimExpr;
+use crate::ir::expr::PrimVar;
use crate::ir::expr::Range;
-use crate::ir::expr::Var;
use crate::sym::analyzer::Analyzer;
use std::ops::Deref;
use tvm_ffi::Any;
@@ -79,7 +84,7 @@ use tvm_ffi::object::ObjectRef;
#[type_final]
pub struct IndexMapObj {
base: Object,
- pub initial_indices: Array<Var>,
+ pub initial_indices: Array<PrimVar>,
pub final_indices: Array<PrimExpr>,
pub inverse_index_map: Option<ObjectRef>,
}
@@ -104,7 +109,7 @@ impl Deref for IndexMap {
impl IndexMapObj {
pub(crate) fn new(
- initial_indices: Array<Var>,
+ initial_indices: Array<PrimVar>,
final_indices: Array<PrimExpr>,
inverse_index_map: Option<ObjectRef>,
) -> Self {
@@ -115,8 +120,8 @@ impl IndexMapObj {
impl IndexMap {
/// Lossless complete-field allocation.
- pub fn new(
- initial_indices: impl Into<Array<Var>>,
+ pub fn from_fields(
+ initial_indices: impl Into<Array<PrimVar>>,
final_indices: impl Into<Array<PrimExpr>>,
inverse_index_map: impl Into<Option<ObjectRef>>,
) -> Self {
@@ -132,9 +137,25 @@ impl IndexMap {
unsafe impl tvm_ffi::object::ImmutableObject for IndexMapObj {}
// tvm-ffi-stubgen(end)
-// The native `tirx.IndexMap` stores its fields as given, as the generated
`IndexMap::new` does.
-
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)`).
+ pub fn new(
+ initial_indices: impl IntoIterator<Item = impl Into<PrimVar>>,
+ final_indices: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
+ inverse_index_map: Option<&IndexMap>,
+ ) -> Result<Self> {
+ let initial_indices =
+
initial_indices.into_iter().map(Into::into).collect::<Array<PrimVar>>();
+ let final_indices = final_indices
+ .into_iter()
+ .map(IntoExpr::into_owned_expr)
+ .collect::<Result<Array<PrimExpr>>>()?;
+ let inverse_index_map = inverse_index_map.map(|map|
ObjectRef::from(map.clone()));
+ Ok(Self::from_fields(initial_indices, final_indices,
inverse_index_map))
+ }
+
/// Return the final indices that `indices` map to
(`tirx.IndexMapMapIndices`).
///
/// # Errors
diff --git a/rust/src/tirx/layout.rs b/rust/src/tirx/layout.rs
index 14a8d3c91..06fb1b23a 100644
--- a/rust/src/tirx/layout.rs
+++ b/rust/src/tirx/layout.rs
@@ -33,7 +33,7 @@
//! use tvm_tirx::tirx::{Axis, Iter, Layout, TileLayout};
//! use tvm_tirx::tvm_ffi::Map;
//!
-//! let m = Axis::get("m")?;
+//! 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");
@@ -58,8 +58,10 @@
// An iterator's extent and stride are `PrimExpr`s, as C++ `IterObj` holds
them.
// tvm-ffi-stubgen(field): tirx.Iter.extent -> PrimExpr
// tvm-ffi-stubgen(field): tirx.Iter.stride -> PrimExpr
+// A tile layout's offsets are `PrimExpr`s, as C++ `TileLayoutObj` holds them.
+// tvm-ffi-stubgen(field): tirx.TileLayout.offset -> Map<Axis, PrimExpr>
+// tvm-ffi-stubgen(import-object): tvm_ffi::Map
// tvm-ffi-stubgen(begin): import-section
-use crate::ir::expr::Expr;
use crate::ir::expr::IntoExpr;
use crate::ir::expr::PrimExpr;
use crate::ir::support;
@@ -124,7 +126,7 @@ impl Layout {
/// Return the number of elements the layout maps, along `axis_name` if
given
/// (`tirx.LayoutGetSize`).
- pub fn size(&self, axis_name: Option<&str>) -> Result<Expr> {
+ pub fn size(&self, axis_name: Option<&str>) -> Result<PrimExpr> {
tvm_ffi::cached_global_func!("tirx.LayoutGetSize")
.call_tuple((self, axis_name.map(String::from)))?
.try_into()
@@ -132,7 +134,7 @@ impl Layout {
/// Return the extent of the coordinates the layout reaches, along
`axis_name` if given
/// (`tirx.LayoutGetSpan`).
- pub fn span(&self, axis_name: Option<&str>) -> Result<Expr> {
+ pub fn span(&self, axis_name: Option<&str>) -> Result<PrimExpr> {
tvm_ffi::cached_global_func!("tirx.LayoutGetSpan")
.call_tuple((self, axis_name.map(String::from)))?
.try_into()
@@ -145,11 +147,13 @@ impl Layout {
/// Returns an error if `coord` does not have one index for each iterator
of the shard.
pub fn apply(
&self,
- coord: impl IntoIterator<Item = impl IntoExpr>,
+ coord: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
analyzer: &Analyzer,
- ) -> Result<Map<String, Expr>> {
- let coord =
-
coord.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
+ ) -> Result<Map<String, PrimExpr>> {
+ let coord = coord
+ .into_iter()
+ .map(IntoExpr::into_owned_expr)
+ .collect::<Result<Array<PrimExpr>>>()?;
tvm_ffi::cached_global_func!("tirx.LayoutApply")
.call_tuple((self, &coord, analyzer))?
.try_into()
@@ -159,9 +163,9 @@ impl Layout {
/// (`tirx.LayoutApplyLinear`).
pub fn apply_linear(
&self,
- coord: impl IntoExpr,
+ coord: impl IntoExpr<PrimExpr>,
analyzer: &Analyzer,
- ) -> Result<Map<String, Expr>> {
+ ) -> Result<Map<String, PrimExpr>> {
tvm_ffi::cached_global_func!("tirx.LayoutApplyLinear")
.call_tuple((self, coord.into_expr()?.borrow(), analyzer))?
.try_into()
@@ -171,14 +175,18 @@ impl Layout {
/// (`tirx.LayoutApplyWithShape`).
pub fn apply_with_shape(
&self,
- coord: impl IntoIterator<Item = impl IntoExpr>,
- shape: impl IntoIterator<Item = impl IntoExpr>,
+ coord: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
+ shape: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
analyzer: &Analyzer,
- ) -> Result<Map<String, Expr>> {
- let coord =
-
coord.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
- let shape =
-
shape.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
+ ) -> Result<Map<String, PrimExpr>> {
+ let coord = coord
+ .into_iter()
+ .map(IntoExpr::into_owned_expr)
+ .collect::<Result<Array<PrimExpr>>>()?;
+ let shape = shape
+ .into_iter()
+ .map(IntoExpr::into_owned_expr)
+ .collect::<Result<Array<PrimExpr>>>()?;
tvm_ffi::cached_global_func!("tirx.LayoutApplyWithShape")
.call_tuple((self, &coord, &shape, analyzer))?
.try_into()
@@ -199,18 +207,18 @@ impl Layout {
pub fn tile(
&self,
outer: impl Into<TileLayout>,
- outer_shape: impl IntoIterator<Item = impl IntoExpr>,
- inner_shape: impl IntoIterator<Item = impl IntoExpr>,
+ outer_shape: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
+ inner_shape: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
) -> Result<Layout> {
let outer = outer.into();
let outer_shape = outer_shape
.into_iter()
.map(IntoExpr::into_owned_expr)
- .collect::<Result<Array<Expr>>>()?;
+ .collect::<Result<Array<PrimExpr>>>()?;
let inner_shape = inner_shape
.into_iter()
.map(IntoExpr::into_owned_expr)
- .collect::<Result<Array<Expr>>>()?;
+ .collect::<Result<Array<PrimExpr>>>()?;
tvm_ffi::cached_global_func!("tirx.LayoutTile")
.call_tuple((self, &outer, &outer_shape, &inner_shape))?
.try_into()
@@ -256,7 +264,7 @@ impl TileLayoutObj {
FieldGetter::new(Self::type_index(), "replica")?.get(self)
}
- pub fn offset(&self) -> Result<Map<Axis, Expr>> {
+ pub fn offset(&self) -> Result<Map<Axis, PrimExpr>> {
FieldGetter::new(Self::type_index(), "offset")?.get(self)
}
}
@@ -272,7 +280,7 @@ impl TileLayout {
pub fn new(
shard: impl IntoIterator<Item = impl Into<Iter>>,
replica: impl IntoIterator<Item = impl Into<Iter>>,
- offset: &Map<Axis, Expr>,
+ offset: &Map<Axis, PrimExpr>,
) -> Self {
let shard = shard.into_iter().map(Into::into).collect::<Array<Iter>>();
let replica =
replica.into_iter().map(Into::into).collect::<Array<Iter>>();
@@ -304,10 +312,12 @@ impl TileLayout {
/// Returns an error if the iterators cannot be split into groups of the
extents of `shape`.
pub fn group(
&self,
- shape: impl IntoIterator<Item = impl IntoExpr>,
+ shape: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
) -> Result<(TileLayout, Array<i64>)> {
- let shape =
-
shape.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
+ let shape = shape
+ .into_iter()
+ .map(IntoExpr::into_owned_expr)
+ .collect::<Result<Array<PrimExpr>>>()?;
let pair: Array<Any> =
tvm_ffi::cached_global_func!("tirx.TileLayoutGroup")
.call_tuple((self, &shape))?
.try_into()?;
@@ -413,8 +423,12 @@ unsafe impl tvm_ffi::object::ImmutableObject for AxisObj {}
impl Axis {
/// Look up the axis named `name`, such as the memory axis `m`,
registering it on first use
/// (`tirx.AxisGet`).
- pub fn get(name: &str) -> Result<Self> {
-
tvm_ffi::cached_global_func!("tirx.AxisGet").call_tuple((String::from(name),))?.try_into()
+ pub fn get(name: &str) -> Self {
+ support::infallible(
+ tvm_ffi::cached_global_func!("tirx.AxisGet")
+ .call_tuple((String::from(name),))
+ .and_then(|any| any.try_into()),
+ )
}
/// Return whether the axis indexes threads (`tirx.AxisIsThreadAxis`).
diff --git a/rust/src/tirx/structural_mutate_with_constraint_context.rs
b/rust/src/tirx/structural_mutate_with_constraint_context.rs
index 758579e94..b405c7da3 100644
--- a/rust/src/tirx/structural_mutate_with_constraint_context.rs
+++ b/rust/src/tirx/structural_mutate_with_constraint_context.rs
@@ -70,7 +70,7 @@ unsafe impl StructuralMutatorBase for
StructuralMutatorWithConstraintContextBase
fn new() -> Result<Self> {
Ok(Self {
base: StructuralMutatorObj::default(),
- constraints: UnsafeCell::new(ConstraintContext::new()?),
+ constraints: UnsafeCell::new(ConstraintContext::new()),
needs_child_context: UnsafeCell::new(false),
})
}
diff --git a/rust/src/tirx/structural_visit_with_constraint_context.rs
b/rust/src/tirx/structural_visit_with_constraint_context.rs
index 5864f3d8c..35a769c27 100644
--- a/rust/src/tirx/structural_visit_with_constraint_context.rs
+++ b/rust/src/tirx/structural_visit_with_constraint_context.rs
@@ -72,7 +72,7 @@ unsafe impl StructuralVisitorBase for
StructuralVisitorWithConstraintContextBase
fn new() -> Result<Self> {
Ok(Self {
base: StructuralVisitorObj::default(),
- constraints: UnsafeCell::new(ConstraintContext::new()?),
+ constraints: UnsafeCell::new(ConstraintContext::new()),
needs_child_context: UnsafeCell::new(false),
})
}
diff --git a/rust/src/tirx/tensor.rs b/rust/src/tirx/tensor.rs
index 27da0d86a..41b914572 100644
--- a/rust/src/tirx/tensor.rs
+++ b/rust/src/tirx/tensor.rs
@@ -36,8 +36,8 @@
//! 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 = Expr::from(PrimVar::new("i", PrimType::int(32), UNKNOWN_LOC));
-//! let point = [Subscript::Index(i), Subscript::Index(0.into())];
+//! 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]");
//!
diff --git a/rust/src/tirx/transform/function_pass.rs
b/rust/src/tirx/transform/function_pass.rs
index db03f802d..673c048bc 100644
--- a/rust/src/tirx/transform/function_pass.rs
+++ b/rust/src/tirx/transform/function_pass.rs
@@ -34,8 +34,9 @@
//! # 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 functions = [(GlobalVar::new("main", UNKNOWN_LOC),
BaseFunc::from(&main))];
-//! # let module = IRModule::from_functions(&functions.into_iter().collect(),
None, &Map::new())?;
+//! # 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)),
diff --git a/rust/src/tirx/transform/transform.rs
b/rust/src/tirx/transform/transform.rs
index 86841f6f2..d11550449 100644
--- a/rust/src/tirx/transform/transform.rs
+++ b/rust/src/tirx/transform/transform.rs
@@ -34,8 +34,9 @@
//! # 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 functions = [(GlobalVar::new("main", UNKNOWN_LOC),
BaseFunc::from(&main))];
-//! # let module = IRModule::from_functions(&functions.into_iter().collect(),
None, &Map::new())?;
+//! # 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()?;
diff --git a/rust/tests/ir/expr.rs b/rust/tests/ir/expr.rs
index eaefbf6b1..9ca6c1390 100644
--- a/rust/tests/ir/expr.rs
+++ b/rust/tests/ir/expr.rs
@@ -95,7 +95,7 @@ fn constants_agree_with_cpp() {
let float_imm = |name: &str, value: f64| {
let native = native(
"ir.FloatImm",
- &[(&parse_dtype(name)).into(), (&value).into(),
(&Location::unknown()).into()],
+ &[(&parse_dtype(name)).into(), (&value).into(),
(&UnknownLoc::new()).into()],
);
let what = format!("FloatImm({name}, {value})");
agree(
@@ -130,7 +130,7 @@ fn constants_agree_with_cpp() {
let int_imm = |name: &str, value: BigInt| {
let native = native(
"ir.IntImm",
- &[(&parse_dtype(name)).into(), (&value).into(),
(&Location::unknown()).into()],
+ &[(&parse_dtype(name)).into(), (&value).into(),
(&UnknownLoc::new()).into()],
);
let what = format!("IntImm({name}, {value})");
agree(
@@ -165,7 +165,7 @@ fn constants_agree_with_cpp() {
let missing: Type = MissingType::new().into();
let generic = |what: &str, ty: &Type, value: Any| {
let native =
- native("ir.GenericConst", &[ty.into(), (&value).into(),
(&Location::unknown()).into()]);
+ native("ir.GenericConst", &[ty.into(), (&value).into(),
(&UnknownLoc::new()).into()]);
agree(what, GenericConst::new(ty, value.clone(), UNKNOWN_LOC), native);
};
let i32_ty_as_type: Type = i32_ty.clone().into();
@@ -177,10 +177,10 @@ fn constants_agree_with_cpp() {
generic("GenericConst(missing type)", &missing, Any::from(i32_ty.clone()));
let int8 = dtype::int(8);
- let native_dtype = native("ir.DataTypeImm", &[(&int8).into(),
(&Location::unknown()).into()]);
+ let native_dtype = native("ir.DataTypeImm", &[(&int8).into(),
(&UnknownLoc::new()).into()]);
agree("DataTypeImm", Ok(DataTypeImm::new(int8, UNKNOWN_LOC)),
native_dtype);
let text = FfiString::from("hi");
- let native_string = native("ir.StringImm", &[(&text).into(),
(&Location::unknown()).into()]);
+ let native_string = native("ir.StringImm", &[(&text).into(),
(&UnknownLoc::new()).into()]);
agree("StringImm", Ok(StringImm::new("hi", UNKNOWN_LOC)), native_string);
}
@@ -192,11 +192,11 @@ fn builtin_constants() {
let int = |dtype: DLDataType, value: i64| {
native(
"ir.IntImm",
- &[(&dtype).into(), (&BigInt::from(value)).into(),
(&Location::unknown()).into()],
+ &[(&dtype).into(), (&BigInt::from(value)).into(),
(&UnknownLoc::new()).into()],
)
};
let float = |dtype: DLDataType, value: f64| {
- native("ir.FloatImm", &[(&dtype).into(), (&value).into(),
(&Location::unknown()).into()])
+ native("ir.FloatImm", &[(&dtype).into(), (&value).into(),
(&UnknownLoc::new()).into()])
};
agree("IntImm::int32", Ok(IntImm::int32(7)), int(dtype::int(32), 7));
agree("IntImm::int32(min)", Ok(IntImm::int32(i32::MIN)),
int(dtype::int(32), i32::MIN as i64));
@@ -240,17 +240,17 @@ fn expressions_agree_with_cpp() {
let i32_ty: Type = PrimType::int(32).into();
let name = FfiString::from("main");
- let native_global = native("ir.GlobalVar", &[(&name).into(),
(&Location::unknown()).into()]);
+ let native_global = native("ir.GlobalVar", &[(&name).into(),
(&UnknownLoc::new()).into()]);
agree_fresh("GlobalVar", Ok(GlobalVar::new("main", UNKNOWN_LOC)),
native_global);
let native_var = native(
"ir.Var",
- &[(&FfiString::from("x")).into(), (&i32_ty).into(),
(&Location::unknown()).into()],
+ &[(&FfiString::from("x")).into(), (&i32_ty).into(),
(&UnknownLoc::new()).into()],
);
agree_fresh("Var", Ok(Var::new("x", i32_ty, UNKNOWN_LOC)), native_var);
let missing: Type = MissingType::new().into();
let native_untyped = native(
"ir.Var",
- &[(&FfiString::from("u")).into(), (&missing).into(),
(&Location::unknown()).into()],
+ &[(&FfiString::from("u")).into(), (&missing).into(),
(&UnknownLoc::new()).into()],
);
agree_fresh("Var(untyped)", Ok(Var::new("u", MissingType::new(),
UNKNOWN_LOC)), native_untyped);
@@ -260,7 +260,7 @@ fn expressions_agree_with_cpp() {
("Tuple(empty)", vec![]),
] {
let array: Array<Expr> = fields.iter().cloned().collect();
- let native_tuple = native("ir.Tuple", &[(&array).into(),
(&Location::unknown()).into()]);
+ let native_tuple = native("ir.Tuple", &[(&array).into(),
(&UnknownLoc::new()).into()]);
agree(what, Ok(Tuple::new(fields, UNKNOWN_LOC)), native_tuple);
}
@@ -273,7 +273,7 @@ fn expressions_agree_with_cpp() {
] {
let native_item = native(
"ir.TupleGetItem",
- &[(&value).into(), (&index).into(), (&Location::unknown()).into()],
+ &[(&value).into(), (&index).into(), (&UnknownLoc::new()).into()],
);
agree(what, TupleGetItem::new(&value, index, UNKNOWN_LOC),
native_item);
}
@@ -283,7 +283,7 @@ fn expressions_agree_with_cpp() {
let array: Array<PrimExpr> = indices.iter().cloned().collect();
let native_load = native(
"ir.TensorLoad",
- &[(&source).into(), (&array).into(),
(&Location::unknown()).into()],
+ &[(&source).into(), (&array).into(), (&UnknownLoc::new()).into()],
);
agree(what, TensorLoad::new(&source, indices, UNKNOWN_LOC),
native_load);
};
@@ -294,7 +294,7 @@ fn expressions_agree_with_cpp() {
let array: Array<Expr> = [text.clone()].into_iter().collect();
let native_load = native(
"ir.TensorLoad",
- &[(&Expr::from(&a)).into(), (&array).into(),
(&Location::unknown()).into()],
+ &[(&Expr::from(&a)).into(), (&array).into(),
(&UnknownLoc::new()).into()],
);
// The type system rejects a string index; its check fails as the native
conversion does.
let rust_load =
@@ -304,7 +304,7 @@ fn expressions_agree_with_cpp() {
let vars: Array<Var> = [Var::from(&x)].into_iter().collect();
let body = op("prim._OpAdd", &x, int32(1));
let native_lambda =
- native("ir.LambdaExpr", &[(&vars).into(), (&body).into(),
(&Location::unknown()).into()]);
+ native("ir.LambdaExpr", &[(&vars).into(), (&body).into(),
(&UnknownLoc::new()).into()]);
agree("LambdaExpr", Ok(LambdaExpr::new([&x], &body, UNKNOWN_LOC)),
native_lambda);
let x_plus_4 = op("prim._OpAdd", &x, int32(4));
@@ -316,16 +316,16 @@ fn expressions_agree_with_cpp() {
("Range(x, n)", PrimExpr::from(&x), PrimExpr::from(&n)),
] {
let native_range =
- native("ir.Range", &[(&begin).into(), (&end).into(),
(&Location::unknown()).into()]);
+ native("ir.Range", &[(&begin).into(), (&end).into(),
(&UnknownLoc::new()).into()]);
agree(what, Range::new(&begin, &end, UNKNOWN_LOC), native_range);
}
// The native constructor spells `[0, n)` as `Range(n)`; Rust gives the
begin.
let native_range =
- native("ir.Range", &[(&n).into(), (&None::<Expr>).into(),
(&Location::unknown()).into()]);
+ native("ir.Range", &[(&n).into(), (&None::<Expr>).into(),
(&UnknownLoc::new()).into()]);
agree("Range(n)", Range::new(int32(0), &n, UNKNOWN_LOC), native_range);
let native_min_extent = native(
"ir.RangeFromMinExtent",
- &[(&Expr::from(&x)).into(), (&int32(4)).into(),
(&Location::unknown()).into()],
+ &[(&Expr::from(&x)).into(), (&int32(4)).into(),
(&UnknownLoc::new()).into()],
);
agree(
"Range::from_min_extent",
@@ -405,7 +405,7 @@ fn calls_agree_with_cpp() {
(&array).into(),
(&attrs).into(),
(&ty_args).into(),
- (&Location::unknown()).into(),
+ (&UnknownLoc::new()).into(),
],
);
let rust = Call::new(ty.as_ref(), &op, args, None, Vec::<Type>::new(),
UNKNOWN_LOC);
diff --git a/rust/tests/ir/location.rs b/rust/tests/ir/location.rs
index 8dc49db9a..4be5cb369 100644
--- a/rust/tests/ir/location.rs
+++ b/rust/tests/ir/location.rs
@@ -26,7 +26,7 @@ use crate::common::{serial, structural_equal};
#[test]
fn source_names() {
let _serial = serial();
- let name = SourceName::get("main.py").unwrap();
+ let name = SourceName::get("main.py");
assert_eq!(name.name.as_str(), "main.py");
- assert!(structural_equal(&name, &SourceName::get("main.py").unwrap()));
+ assert!(structural_equal(&name, &SourceName::get("main.py")));
}
diff --git a/rust/tests/ir/main.rs b/rust/tests/ir/main.rs
index ce15f054c..1c6c99bf5 100644
--- a/rust/tests/ir/main.rs
+++ b/rust/tests/ir/main.rs
@@ -57,7 +57,7 @@ fn int32_var(name: &str) -> PrimVar {
}
fn op(name: &str, a: impl Into<Expr>, b: impl Into<Expr>) -> PrimExpr {
- global(name).call_tuple((a.into(), b.into(),
Location::unknown())).unwrap().try_into().unwrap()
+ global(name).call_tuple((a.into(), b.into(),
UnknownLoc::new())).unwrap().try_into().unwrap()
}
/// A memory scope registered once for all the tests: Rust has no lookup of
the builtin ones.
@@ -88,7 +88,7 @@ fn tensor_var(name: &str) -> Var {
AnyView::from(&scope),
AnyView::from(&alignment),
AnyView::from(&layout),
- AnyView::from(&Location::unknown()),
+ AnyView::from(&UnknownLoc::new()),
])
.unwrap()
.try_into()
@@ -107,7 +107,7 @@ fn function(value: i64) -> BaseFunc {
let ret_type: Type = PrimType::int(32).into();
let attrs = DictAttrs::new(Map::new());
global("tirx.Function")
- .call_tuple((params, body, ret_type, attrs, Location::unknown()))
+ .call_tuple((params, body, ret_type, attrs, UnknownLoc::new()))
.unwrap()
.try_into()
.unwrap()
diff --git a/rust/tests/ir/module.rs b/rust/tests/ir/module.rs
index 0642a8849..4c1272202 100644
--- a/rust/tests/ir/module.rs
+++ b/rust/tests/ir/module.rs
@@ -31,7 +31,7 @@ fn modules() {
let func = function(0);
let main = GlobalVar::new("main", UNKNOWN_LOC);
let functions: Map<GlobalVar, BaseFunc> = [(main.clone(),
func.clone())].into_iter().collect();
- let module = IRModule::from_functions(&functions, None,
&Map::new()).unwrap();
+ let module = IRModule::new(&functions, DictAttrs::new(Map::new()),
&Map::new()).unwrap();
assert!(module.contains("main").unwrap());
assert!(module.contains_global_var(&main).unwrap());
assert!(module.has_global_var("main").unwrap());
@@ -53,7 +53,8 @@ fn modules() {
let extra_var = GlobalVar::new("extra", UNKNOWN_LOC);
let extra_functions: Map<GlobalVar, BaseFunc> =
[(extra_var.clone(), function(4))].into_iter().collect();
- module.update(&IRModule::from_functions(&extra_functions, None,
&Map::new()).unwrap()).unwrap();
+ let extra = IRModule::new(&extra_functions, DictAttrs::new(Map::new()),
&Map::new()).unwrap();
+ module.update(&extra).unwrap();
assert!(module.contains("extra").unwrap());
module.update_global_info("info", Vec::<GlobalInfo>::new()).unwrap();
assert!(module.global_infos.get(&FfiString::from("info")).is_ok());
@@ -101,7 +102,8 @@ fn modules_agree_with_cpp() {
let attrs_arg = attrs.cloned();
let native_module =
native("ir.IRModule", &[(&functions).into(), (&attrs_arg).into(),
(&infos).into()]);
- let rust = IRModule::from_functions(&functions, attrs, &infos);
+ let attrs = attrs.cloned().unwrap_or_else(||
DictAttrs::new(Map::new()));
+ let rust = IRModule::new(&functions, attrs, &infos);
if let (Ok(rust), Ok(native_module)) = (&rust, &native_module) {
let native_module: IRModule =
native_module.clone().try_into().unwrap();
assert!(
diff --git a/rust/tests/ir/op_def.rs b/rust/tests/ir/op_def.rs
index 994e3f158..94de9e05e 100644
--- a/rust/tests/ir/op_def.rs
+++ b/rust/tests/ir/op_def.rs
@@ -147,7 +147,7 @@ fn ops_registered_from_rust_are_used_natively() {
Array::<Expr>::from_vec(vec![x.clone().into()]),
Option::<Attrs>::None,
Array::<Type>::from_vec(vec![]),
- tvm_tirx::ir::Location::unknown(),
+ tvm_tirx::ir::UnknownLoc::new(),
))
.unwrap()
.try_into()
diff --git a/rust/tests/ir/prim/expr.rs b/rust/tests/ir/prim/expr.rs
index 41dad188e..42cd89c79 100644
--- a/rust/tests/ir/prim/expr.rs
+++ b/rust/tests/ir/prim/expr.rs
@@ -21,7 +21,7 @@
use tvm_tirx::ir::prim;
use tvm_tirx::ir::{
- UNKNOWN_LOC, Location, MissingType, Expr, PrimExpr, PrimType, PrimVar,
StringImm, Type, Var,
+ UNKNOWN_LOC, UnknownLoc, MissingType, Expr, PrimExpr, PrimType, PrimVar,
StringImm, Type, Var,
};
use tvm_tirx::tvm_ffi::{Any, AnyCompatible, AnyView, Result};
@@ -84,7 +84,7 @@ fn node_constructors() {
/// The native constructor `name` on `args` and the location.
fn native_with_loc<T: AnyCompatible>(name: &str, args: &[&T]) -> Result<Any> {
- let loc = Location::unknown();
+ let loc = UnknownLoc::new();
let mut views: Vec<AnyView> = args.iter().map(|arg|
AnyView::from(*arg)).collect();
views.push(AnyView::from(&loc));
native(name, &views)
@@ -174,10 +174,8 @@ fn other_nodes_agree_with_cpp() {
("Cast(int32 to float32x4)",
PrimType::float(32).with_lanes(4).unwrap(), &x),
] {
let target_ty: Type = target.clone().into();
- let native_cast = native(
- "prim.Cast",
- &[(&target_ty).into(), value.into(),
(&Location::unknown()).into()],
- );
+ let native_cast =
+ native("prim.Cast", &[(&target_ty).into(), value.into(),
(&UnknownLoc::new()).into()]);
agree(what, prim::Cast::new(target, value, UNKNOWN_LOC), native_cast);
}
@@ -207,7 +205,7 @@ fn other_nodes_agree_with_cpp() {
let body = PrimExpr::from(if what == "Let" { &z } else { &x });
let native_let = native(
"prim.Let",
- &[var.into(), value.into(), (&body).into(),
(&Location::unknown()).into()],
+ &[var.into(), value.into(), (&body).into(),
(&UnknownLoc::new()).into()],
);
agree(what, prim::Let::new(var, value, &body, UNKNOWN_LOC),
native_let);
}
diff --git a/rust/tests/ir/prim/op.rs b/rust/tests/ir/prim/op.rs
index 6f2316e00..123856889 100644
--- a/rust/tests/ir/prim/op.rs
+++ b/rust/tests/ir/prim/op.rs
@@ -26,8 +26,8 @@
use tvm_tirx::ir::prim::op::{self, ExprResult, PrimOps};
use tvm_tirx::ir::prim::{self, Add, Broadcast, BroadcastObj, Cast, CastObj,
Not, Ramp, RampObj, Sub};
use tvm_tirx::ir::{
- Call, Expr, FloatImm, IntImm, Location, Op, PrimExpr, PrimType, Range,
SourceLoc, SourceName,
- StringImm, Type, UnknownLocObj, UNKNOWN_LOC,
+ Call, Expr, FloatImm, IntImm, Location, UnknownLoc, Op, PrimExpr,
PrimType, Range, SourceLoc,
+ SourceName, StringImm, Type, UnknownLocObj, UNKNOWN_LOC,
};
use tvm_tirx::sym;
use tvm_tirx::tvm_ffi::{
@@ -73,12 +73,12 @@ where
/// The native operator `name` applied to two expressions.
fn native2(name: &str, a: impl Into<Expr>, b: impl Into<Expr>) -> Expr {
- global(name).call_tuple((a.into(), b.into(),
Location::unknown())).unwrap().try_into().unwrap()
+ global(name).call_tuple((a.into(), b.into(),
UnknownLoc::new())).unwrap().try_into().unwrap()
}
/// The native unary function `name`.
fn native1(name: &str, x: impl Into<Expr>) -> Result<Any> {
- global(name).call_tuple((x.into(), Location::unknown()))
+ global(name).call_tuple((x.into(), UnknownLoc::new()))
}
fn same(built: ExprResult, expected: &impl AnyCompatible) -> bool {
@@ -262,7 +262,7 @@ fn binary_operators_agree_with_cpp() {
let mut outcome = Outcome::default();
for a in &operands {
for b in &operands {
- let native = native_op.call_tuple((a.clone(), b.clone(),
Location::unknown()));
+ let native = native_op.call_tuple((a.clone(), b.clone(),
UnknownLoc::new()));
let what = || format!("{name}({}, {})", repr(a), repr(b));
outcome.check(what, operator(a, b).try_into(), native);
}
@@ -284,7 +284,7 @@ fn unary_and_ternary_operators_agree_with_cpp() {
// own check rejects the same operands with another message.
let bitwise_not = global("prim.BitwiseNot");
for a in &operands {
- let native = bitwise_not.call_tuple((a.clone(), Location::unknown()));
+ let native = bitwise_not.call_tuple((a.clone(), UnknownLoc::new()));
match a.bitwise_not().try_into() {
Ok(rust) => outcome.check(|| format!("~{}", repr(a)), Ok(rust),
native),
Err(_) => assert!(native.is_err(), "~{}: only Rust rejects it",
repr(a)),
@@ -307,7 +307,7 @@ fn unary_and_ternary_operators_agree_with_cpp() {
for t in &values {
for f in &values {
let native =
- if_then_else.call_tuple((c.clone(), t.clone(), f.clone(),
Location::unknown()));
+ if_then_else.call_tuple((c.clone(), t.clone(), f.clone(),
UnknownLoc::new()));
let what = || format!("if_then_else({}, {}, {})", repr(c),
repr(t), repr(f));
outcome.check(what, c.if_then_else(t, f).try_into(), native);
}
@@ -326,26 +326,26 @@ fn literals_and_compound_expressions_agree_with_cpp() {
outcome.check(
|| "x + 3".into(),
(&x + 3).try_into(),
- add.call_tuple((x.clone(), imm("int64", 3), Location::unknown())),
+ add.call_tuple((x.clone(), imm("int64", 3), UnknownLoc::new())),
);
outcome.check(
|| "3 + x".into(),
(3i32 + &x).try_into(),
- add.call_tuple((imm("int64", 3), x.clone(), Location::unknown())),
+ add.call_tuple((imm("int64", 3), x.clone(), UnknownLoc::new())),
);
let floordiv = global("prim._OpFloorDiv");
let rust = ((&x + 0) * 1).floordiv(1);
outcome.check(
|| "((x + 0) * 1) // 1".into(),
rust.try_into(),
- floordiv.call_tuple((x.clone(), imm("int64", 1), Location::unknown())),
+ floordiv.call_tuple((x.clone(), imm("int64", 1), UnknownLoc::new())),
);
// Unary minus folds a constant whose negation fits its type and is `0 -
x` otherwise.
let sub = global("prim._OpSub");
for value in [imm("int64", i64::MIN), imm("int32", 5), fimm("float32",
2.5), prim_var("int32")]
{
let zero = op::make_const(value.ty(), 0).unwrap();
- let native = sub.call_tuple((zero, value.clone(),
Location::unknown()));
+ let native = sub.call_tuple((zero, value.clone(), UnknownLoc::new()));
outcome.check(|| format!("-{}", repr(&value)), (-&value).try_into(),
native);
}
outcome.assert_agrees("literals");
@@ -429,7 +429,7 @@ fn
expression_results_pass_into_functions_and_constructors() {
fn located_variants() {
let _serial = serial();
let x = var("x", PrimType::int(32));
- let here: Location = SourceLoc::new(SourceName::get("ops.rs").unwrap(), 1,
2, 3, 4).into();
+ let here: Location = SourceLoc::new(SourceName::get("ops.rs"), 1, 2, 3,
4).into();
let located: prim::FloorDiv = op::floordiv_at(&x, 4,
here.clone()).unwrap().downcast().unwrap();
assert!(located.loc.same_as(&here));
let unlocated: prim::FloorDiv = op::floordiv(&x,
4).unwrap().downcast().unwrap();
@@ -474,7 +474,7 @@ fn select_folding_and_the_first_error() {
let cond = native2("prim._OpGE", &r, imm("int32", 0));
let minus = native2("prim._OpSub", &r, imm("int32", 1));
let expected: Expr = global("prim.Select")
- .call_tuple((cond, Expr::from(&r), minus, Location::unknown()))
+ .call_tuple((cond, Expr::from(&r), minus, UnknownLoc::new()))
.unwrap()
.try_into()
.unwrap();
@@ -489,11 +489,8 @@ fn select_folding_and_the_first_error() {
let Err(error) = PrimExpr::try_from((f.truncdiv(&x) + 1).less(&x)) else {
panic!("expected an error")
};
- let direct = global("prim._OpTruncDiv").call_tuple((
- Expr::from(&f),
- Expr::from(&x),
- Location::unknown(),
- ));
+ let direct =
+ global("prim._OpTruncDiv").call_tuple((Expr::from(&f), Expr::from(&x),
UnknownLoc::new()));
let direct = direct.err().expect("the native operation raises");
assert_eq!(error.kind(), direct.kind());
assert_eq!(error.message(), direct.message());
@@ -592,8 +589,7 @@ fn limits_agree_with_cpp() {
];
for name in types {
let ty = prim(name);
- let native =
- |function: &str| global(function).call_tuple((ty.clone(),
Location::unknown()));
+ let native = |function: &str| global(function).call_tuple((ty.clone(),
UnknownLoc::new()));
agree(&format!("infinity({name})"), op::infinity(&ty),
native("prim.infinity"));
agree(&format!("max_value({name})"), op::max_value(&ty),
native("prim.max_value"));
agree(&format!("min_value({name})"), op::min_value(&ty),
native("prim.min_value"));
@@ -627,7 +623,7 @@ fn cast_agrees_with_cpp() {
for name in targets {
let ty = prim(name);
let native =
- global("prim._cast").call_tuple((x.clone(), ty.clone(),
Location::unknown()));
+ global("prim._cast").call_tuple((x.clone(), ty.clone(),
UnknownLoc::new()));
agree(&format!("cast({}, {name})", repr(x)),
PrimExpr::try_from(x.cast(ty)), native);
}
}
@@ -664,12 +660,12 @@ fn type_matching_agrees_with_the_native_operators() {
let (x, y) = (var("x", prim(a)), var("y", prim(b)));
// The raw `Add` of the matched sides is what C++ `x + y` builds when
nothing folds.
let native =
- global("prim._OpAdd").call_tuple((Expr::from(&x), Expr::from(&y),
Location::unknown()));
+ global("prim._OpAdd").call_tuple((Expr::from(&x), Expr::from(&y),
UnknownLoc::new()));
agree(&format!("{a} + {b}"), raw_add(x.into(), y.into()), native);
}
// A constant side is cast by folding.
let (x, three) = (prim_var("float32"), imm("int32", 3));
- let native = global("prim._OpAdd").call_tuple((x.clone(), three.clone(),
Location::unknown()));
+ let native = global("prim._OpAdd").call_tuple((x.clone(), three.clone(),
UnknownLoc::new()));
agree("float32 + 3", raw_add(x, three), native);
}
@@ -685,7 +681,7 @@ fn pow_and_fmod_agree_with_cpp() {
(prim_var("int32"), var("j", PrimType::int(32)).into()),
];
for (x, y) in &pairs {
- let native = global("prim._OpPow").call_tuple((x.clone(), y.clone(),
Location::unknown()));
+ let native = global("prim._OpPow").call_tuple((x.clone(), y.clone(),
UnknownLoc::new()));
agree(&format!("pow({}, {})", repr(x), repr(y)),
PrimExpr::try_from(x.pow(y)), native);
}
// `fmod` has no global: the sides are matched and the call is typed as
them.
@@ -754,7 +750,7 @@ fn intrinsics_are_typed_calls() {
/// C++ `cast(value, t)` through its global.
fn native_cast(value: impl Into<Expr>, name: &str) -> Expr {
global("prim._cast")
- .call_tuple((value.into(), prim(name), Location::unknown()))
+ .call_tuple((value.into(), prim(name), UnknownLoc::new()))
.unwrap()
.try_into()
.unwrap()
diff --git a/rust/tests/ir/prim/vector_expr.rs
b/rust/tests/ir/prim/vector_expr.rs
index 487b9e158..ab716bf41 100644
--- a/rust/tests/ir/prim/vector_expr.rs
+++ b/rust/tests/ir/prim/vector_expr.rs
@@ -20,7 +20,7 @@
//! Tests of `ir::prim::vector_expr`.
use tvm_tirx::ir::prim;
-use tvm_tirx::ir::{UNKNOWN_LOC, Location, Call, Expr, IntImm, Op, PrimExpr,
PrimType, Type};
+use tvm_tirx::ir::{UNKNOWN_LOC, UnknownLoc, Call, Expr, IntImm, Op, PrimExpr,
PrimType, Type};
use tvm_tirx::tvm_ffi::Array;
use super::{dtype, var};
@@ -96,7 +96,7 @@ fn vector_nodes_agree_with_cpp() {
for (what, base, stride, lanes) in cases {
let native_ramp = native(
"prim.Ramp",
- &[(&base).into(), (&stride).into(), (&lanes).into(),
(&Location::unknown()).into()],
+ &[(&base).into(), (&stride).into(), (&lanes).into(),
(&UnknownLoc::new()).into()],
);
agree(what, prim::Ramp::new(&base, &stride, &lanes, UNKNOWN_LOC),
native_ramp);
}
@@ -109,7 +109,7 @@ fn vector_nodes_agree_with_cpp() {
] {
let native_broadcast = native(
"prim.Broadcast",
- &[(&value).into(), (&lanes).into(), (&Location::unknown()).into()],
+ &[(&value).into(), (&lanes).into(), (&UnknownLoc::new()).into()],
);
agree(what, prim::Broadcast::new(&value, &lanes, UNKNOWN_LOC),
native_broadcast);
}
@@ -134,7 +134,7 @@ fn vector_nodes_agree_with_cpp() {
(vectors.iter().cloned().collect(),
indices.iter().cloned().collect());
let native_shuffle = native(
"prim.Shuffle",
- &[(&vector_array).into(), (&index_array).into(),
(&Location::unknown()).into()],
+ &[(&vector_array).into(), (&index_array).into(),
(&UnknownLoc::new()).into()],
);
agree(what, prim::Shuffle::new(vectors, indices, UNKNOWN_LOC),
native_shuffle);
}
diff --git a/rust/tests/ir/stmt.rs b/rust/tests/ir/stmt.rs
index 942bf0e97..c07cf91a4 100644
--- a/rust/tests/ir/stmt.rs
+++ b/rust/tests/ir/stmt.rs
@@ -55,7 +55,7 @@ fn statements() {
UNKNOWN_LOC,
)
.unwrap();
- assert_eq!(for_.kind, kind as i32);
+ assert_eq!(for_.kind, kind);
assert_eq!(for_.body.seq.len(), 1);
let less = op("prim._OpLT", &x, int32(4));
let while_ = While::new(&less, &body, UNKNOWN_LOC).unwrap();
@@ -110,7 +110,7 @@ fn statements_agree_with_cpp() {
for (what, var, value) in [("Bind", &y, Expr::from(&x)), ("Bind(type)",
&y, Expr::from(&f))] {
let native_bind =
- native("ir.Bind", &[var.into(), (&value).into(),
(&Location::unknown()).into()]);
+ native("ir.Bind", &[var.into(), (&value).into(),
(&UnknownLoc::new()).into()]);
agree(what, Bind::new(var, &value, UNKNOWN_LOC), native_bind);
}
@@ -121,7 +121,7 @@ fn statements_agree_with_cpp() {
{
let native_assert = native(
"ir.AssertStmt",
- &[(&condition).into(), (&kind).into(), (&parts).into(),
(&Location::unknown()).into()],
+ &[(&condition).into(), (&kind).into(), (&parts).into(),
(&UnknownLoc::new()).into()],
);
agree(what, AssertStmt::new(&condition, &kind, parts.iter(),
UNKNOWN_LOC), native_assert);
}
@@ -192,7 +192,7 @@ fn statements_agree_with_cpp() {
),
];
for (what, var, begin, end, step, body, annotations) in cases {
- let kind = ForKind::Parallel as i32;
+ let kind = ForKind::Parallel.as_raw();
let native_for = native(
"ir.For",
&[
@@ -203,7 +203,7 @@ fn statements_agree_with_cpp() {
(&kind).into(),
(&body).into(),
(&annotations).into(),
- (&Location::unknown()).into(),
+ (&UnknownLoc::new()).into(),
],
);
let rust = For::new(
@@ -225,7 +225,7 @@ fn statements_agree_with_cpp() {
{
let native_while = native(
"ir.While",
- &[(&condition).into(), (&evaluate).into(),
(&Location::unknown()).into()],
+ &[(&condition).into(), (&evaluate).into(),
(&UnknownLoc::new()).into()],
);
agree(what, While::new(&condition, &evaluate, UNKNOWN_LOC),
native_while);
}
@@ -242,7 +242,7 @@ fn statements_agree_with_cpp() {
let array: Array<PrimExpr> = indices.iter().cloned().collect();
let native_store = native(
"ir.TensorStore",
- &[(&dest).into(), (&array).into(), (&value).into(),
(&Location::unknown()).into()],
+ &[(&dest).into(), (&array).into(), (&value).into(),
(&UnknownLoc::new()).into()],
);
agree(what, TensorStore::new(&dest, indices, &value, UNKNOWN_LOC),
native_store);
}
@@ -256,7 +256,7 @@ fn statements_agree_with_cpp() {
];
for (what, stmts) in sequences {
let array: Array<Stmt> = stmts.iter().cloned().collect();
- let native_seq = native("ir.SeqStmt", &[(&array).into(),
(&Location::unknown()).into()]);
+ let native_seq = native("ir.SeqStmt", &[(&array).into(),
(&UnknownLoc::new()).into()]);
agree(what, Ok(SeqStmt::new(stmts, UNKNOWN_LOC)), native_seq);
}
@@ -272,17 +272,17 @@ fn statements_agree_with_cpp() {
(&condition).into(),
(&then_case).into(),
(&else_case).into(),
- (&Location::unknown()).into(),
+ (&UnknownLoc::new()).into(),
],
);
agree(what, If::new(&condition, &then_case, else_case.as_ref(),
UNKNOWN_LOC), native_if);
}
let native_evaluate =
- native("ir.Evaluate", &[(&Expr::from(&x)).into(),
(&Location::unknown()).into()]);
+ native("ir.Evaluate", &[(&Expr::from(&x)).into(),
(&UnknownLoc::new()).into()]);
agree("Evaluate", Ok(Evaluate::new(&x, UNKNOWN_LOC)), native_evaluate);
let native_return =
- native("ir.Return", &[(&Expr::from(&x)).into(),
(&Location::unknown()).into()]);
+ native("ir.Return", &[(&Expr::from(&x)).into(),
(&UnknownLoc::new()).into()]);
agree("Return", Ok(Return::new(&x, UNKNOWN_LOC)), native_return);
}
@@ -373,7 +373,7 @@ fn regions_agree_with_cpp() {
attrs.into(),
(&body).into(),
(&results_array).into(),
- (&Location::unknown()).into(),
+ (&UnknownLoc::new()).into(),
],
);
let rust = RegionStmt::new(op, args, params, attrs, &body, results,
UNKNOWN_LOC);
diff --git a/rust/tests/ir/transform/core.rs b/rust/tests/ir/transform/core.rs
index 496c4b89a..862480fbe 100644
--- a/rust/tests/ir/transform/core.rs
+++ b/rust/tests/ir/transform/core.rs
@@ -29,12 +29,12 @@ use crate::common::{empty_module, serial};
fn pass_context_scope() {
let _serial = serial();
let default_level = PassContext::current().unwrap().opt_level;
- let ctx = PassContext::create(3, ["a"], ["b"], no_instruments(),
None).unwrap();
+ let ctx = PassContext::new(3, ["a"], ["b"], no_instruments(),
None).unwrap();
assert_eq!(ctx.required_pass.len(), 1);
{
let _outer = ctx.enter().unwrap();
assert_eq!(PassContext::current().unwrap().opt_level, 3);
- let inner = PassContext::create(1, [""; 0], [""; 0], no_instruments(),
None).unwrap();
+ let inner = PassContext::new(1, [""; 0], [""; 0], no_instruments(),
None).unwrap();
{
let _inner = inner.enter().unwrap();
assert_eq!(PassContext::current().unwrap().opt_level, 1);
@@ -47,7 +47,7 @@ fn pass_context_scope() {
assert!(!configs.is_empty());
let unknown: Map<FfiString, Any> =
[(FfiString::from("no.such.option"),
Any::from(1i64))].into_iter().collect();
- assert!(PassContext::create(2, [""; 0], [""; 0], no_instruments(),
Some(&unknown)).is_err());
+ assert!(PassContext::new(2, [""; 0], [""; 0], no_instruments(),
Some(&unknown)).is_err());
}
#[test]
@@ -78,12 +78,11 @@ fn module_and_sequential_passes() {
print_ir("between").unwrap(),
logging_pass("b", &log).into(),
],
- 0,
- "seq",
+ PassInfo::new(0, "seq"),
);
assert_eq!(seq.passes().unwrap().len(), 3);
assert_eq!(Pass::from(&seq).info().unwrap().name.as_str(), "seq");
- let ctx = PassContext::create(2, [""; 0], ["b"], no_instruments(),
None).unwrap();
+ let ctx = PassContext::new(2, [""; 0], ["b"], no_instruments(),
None).unwrap();
{
let _scope = ctx.enter().unwrap();
Pass::from(&seq).run(empty_module()).unwrap();
diff --git a/rust/tests/ir/transform/instrument.rs
b/rust/tests/ir/transform/instrument.rs
index b2f3b6248..7b017f5d4 100644
--- a/rust/tests/ir/transform/instrument.rs
+++ b/rust/tests/ir/transform/instrument.rs
@@ -58,7 +58,7 @@ fn instruments() {
assert_eq!(probe.name().unwrap().as_str(), "probe");
let silent = BasePassInstrument::new("silent",
PassInstrumentHooks::default());
- let ctx = PassContext::create(2, [""; 0], [""; 0], [&probe],
None).unwrap();
+ let ctx = PassContext::new(2, [""; 0], [""; 0], [&probe], None).unwrap();
{
let _scope = ctx.enter().unwrap();
Pass::from(logging_pass("kept", &log)).run(empty_module()).unwrap();
@@ -77,7 +77,7 @@ fn pass_timing() {
let _serial = serial();
let log = Log::default();
let timing = pass_timing_instrument().unwrap();
- let ctx = PassContext::create(2, [""; 0], [""; 0], [timing],
None).unwrap();
+ let ctx = PassContext::new(2, [""; 0], [""; 0], [timing], None).unwrap();
let _scope = ctx.enter().unwrap();
Pass::from(logging_pass("timed", &log)).run(empty_module()).unwrap();
let profile = render_pass_timing().unwrap();
diff --git a/rust/tests/ir/type.rs b/rust/tests/ir/type.rs
index 5754dc922..5418f71ec 100644
--- a/rust/tests/ir/type.rs
+++ b/rust/tests/ir/type.rs
@@ -24,8 +24,9 @@
use std::mem::offset_of;
use tvm_tirx::ir::{
- UNKNOWN_LOC, Location, dtype, Evaluate, FuncType, MissingType, NoneType,
PrimType, PtrType,
- TensorRegionType, TupleType, Type, TypeObj, Var, VarObj,
+ UNKNOWN_LOC, Location, UnknownLoc, dtype, Evaluate, FuncType, MissingType,
NoneType,
+ OpaqueType, PrimType, PtrType, StringType, TensorRegionType, TupleType,
Type, TypeObj, Var,
+ VarObj,
};
use tvm_tirx::tvm_ffi::{def_object_ref, rs_object, structural_map};
use tvm_tirx::tvm_ffi::{
@@ -101,13 +102,14 @@ fn types_agree_with_cpp() {
[("TupleType", vec![int32.clone(), float32.clone()]),
("TupleType(empty)", vec![])]
{
let array: Array<Type> = fields.iter().cloned().collect();
- let native_tuple =
- native("ir.TupleType", &[(&array).into(),
(&Location::unknown()).into()]);
+ let native_tuple = native("ir.TupleType", &[(&array).into(),
(&UnknownLoc::new()).into()]);
agree(what, Ok(TupleType::new(fields, UNKNOWN_LOC)), native_tuple);
}
let args: Array<Type> = [int32.clone()].into_iter().collect();
let native_func = native("ir.FuncType", &[(&args).into(),
(&float32).into()]);
agree("FuncType", Ok(FuncType::new([&int32], &float32)), native_func);
+ agree("OpaqueType", Ok(OpaqueType::new()), native("ir.OpaqueType", &[]));
+ agree("StringType", Ok(StringType::new()), native("ir.StringType", &[]));
let none: Type = native("ir.NoneType", &[]).unwrap().try_into().unwrap();
assert!(NoneType::new().same_as(&none));
@@ -141,7 +143,7 @@ impl PairType {
}
fn pair(lhs: Type, rhs: Type, tag: i64, note: &str) -> PairType {
- PairTypeObj::new(TypeObj::new(Location::unknown()), lhs, rhs, tag,
note).into()
+ PairTypeObj::new(TypeObj::new(UNKNOWN_LOC.into()), lhs, rhs, tag,
note).into()
}
fn int_float_pair(tag: i64, note: &str) -> PairType {
diff --git a/rust/tests/sym/int_set.rs b/rust/tests/sym/int_set.rs
index 7263f5254..2dfb1c7eb 100644
--- a/rust/tests/sym/int_set.rs
+++ b/rust/tests/sym/int_set.rs
@@ -19,7 +19,7 @@
//! Tests of `sym::int_set`: integer sets and region bounds.
-use tvm_tirx::ir::{IntImm, Location, PrimExpr, Range, Var};
+use tvm_tirx::ir::{IntImm, UnknownLoc, PrimExpr, Range, Var};
use tvm_tirx::sym::*;
use tvm_tirx::tvm_ffi::{Map, Object, ObjectRefCore};
@@ -93,7 +93,7 @@ fn int_sets() {
fn region_bounds() {
let x = var("x");
let var_dom = Map::from_iter([(Var::from(&x), range(0, 16))]);
- let region = [Range::from_fields(x.clone(), int32(1),
Location::unknown())];
+ let region = [Range::from_fields(x.clone(), int32(1), UnknownLoc::new())];
let predicate = PrimExpr::from(true);
let lower = estimate_region_lower_bound(region.clone(), &var_dom,
&predicate, None).unwrap();
diff --git a/rust/tests/sym/iter_affine_map.rs
b/rust/tests/sym/iter_affine_map.rs
index fbd42f2d3..384be7d5e 100644
--- a/rust/tests/sym/iter_affine_map.rs
+++ b/rust/tests/sym/iter_affine_map.rs
@@ -90,7 +90,7 @@ fn iter_maps_decompose_into_marks_and_splits() {
fn iter_maps() {
let analyzer = Analyzer::new();
let (x, y) = (var("x"), var("y"));
- let dom = Map::from_iter([(Var::from(&x), range(0, 4)), (Var::from(&y),
range(0, 4))]);
+ let dom = Map::from_iter([(x.clone(), range(0, 4)), (y.clone(), range(0,
4))]);
let index = op("prim._OpAdd", op("prim._OpMul", &x, int32(4)), &y);
let predicate = PrimExpr::from(true);
diff --git a/rust/tests/sym/main.rs b/rust/tests/sym/main.rs
index 9025aad71..3615d6fcc 100644
--- a/rust/tests/sym/main.rs
+++ b/rust/tests/sym/main.rs
@@ -27,7 +27,7 @@ mod analyzer;
mod int_set;
mod iter_affine_map;
-use tvm_tirx::ir::{UNKNOWN_LOC, Expr, IntImm, Location, PrimExpr, PrimType,
PrimVar, Range, Var};
+use tvm_tirx::ir::{UNKNOWN_LOC, Expr, IntImm, UnknownLoc, PrimExpr, PrimType,
PrimVar, Range, Var};
use tvm_tirx::sym::{Analyzer, ConstraintContext};
use tvm_tirx::tvm_ffi::String as FfiString;
use tvm_tirx::tvm_ffi::object::ObjectRef;
@@ -39,19 +39,19 @@ fn int32(value: i64) -> PrimExpr {
}
fn op(name: &str, a: impl Into<Expr>, b: impl Into<Expr>) -> PrimExpr {
- global(name).call_tuple((a.into(), b.into(),
Location::unknown())).unwrap().try_into().unwrap()
+ global(name).call_tuple((a.into(), b.into(),
UnknownLoc::new())).unwrap().try_into().unwrap()
}
fn var(name: &str) -> PrimVar {
global("ir.Var")
- .call_tuple((FfiString::from(name), PrimType::int(32),
Location::unknown()))
+ .call_tuple((FfiString::from(name), PrimType::int(32),
UnknownLoc::new()))
.unwrap()
.try_into()
.unwrap()
}
fn range(begin: i64, extent: i64) -> Range {
- Range::from_fields(int32(begin), int32(extent), Location::unknown())
+ Range::from_fields(int32(begin), int32(extent), UnknownLoc::new())
}
fn context(constraints: impl IntoIterator<Item = ObjectRef>) ->
ConstraintContext {
diff --git a/rust/tests/tirx/analysis/analysis.rs
b/rust/tests/tirx/analysis/analysis.rs
index 672638670..71494ff1c 100644
--- a/rust/tests/tirx/analysis/analysis.rs
+++ b/rust/tests/tirx/analysis/analysis.rs
@@ -24,7 +24,7 @@ use tvm_tirx::ir::PrimType;
use tvm_tirx::tirx::analysis;
use crate::common::{serial, structural_equal};
-use crate::{as_i64, int32, module_of, op, var, void_function};
+use crate::{as_i64, int32, module_of, op, prim_var, var, void_function};
#[test]
fn well_formedness() {
@@ -38,7 +38,7 @@ fn well_formedness() {
#[test]
fn expression_analyses() {
let _serial = serial();
- let (x, y) = (var("x", PrimType::int(32)), var("y", PrimType::int(32)));
+ let (x, y) = (prim_var("x", PrimType::int(32)), prim_var("y",
PrimType::int(32)));
let undefined = analysis::undefined_vars(op("prim._OpAdd", &x, &y),
[x.clone()]).unwrap();
assert_eq!(undefined.len(), 1);
assert!(structural_equal(&undefined.get(0).unwrap(), &y));
diff --git a/rust/tests/tirx/function.rs b/rust/tests/tirx/function.rs
index 794f3eb2d..b829c51d3 100644
--- a/rust/tests/tirx/function.rs
+++ b/rust/tests/tirx/function.rs
@@ -20,7 +20,7 @@
//! Tests of `tirx::function`.
use tvm_tirx::ir::{
- UNKNOWN_LOC, Location, DictAttrs, MissingType, PrimType, Return, SeqStmt,
Type, Var,
+ UNKNOWN_LOC, UnknownLoc, DictAttrs, MissingType, PrimType, Return,
SeqStmt, Type, Var,
};
use tvm_tirx::tirx::{renew_def, Function};
use tvm_tirx::tvm_ffi::{Any, Array, Map, String as FfiString};
@@ -67,7 +67,7 @@ fn functions_agree_with_cpp() {
(&body).into(),
ret_type.into(),
(&attrs).into(),
- (&Location::unknown()).into(),
+ (&UnknownLoc::new()).into(),
],
);
agree(
diff --git a/rust/tests/tirx/index_map.rs b/rust/tests/tirx/index_map.rs
index d73b92c52..f92320b37 100644
--- a/rust/tests/tirx/index_map.rs
+++ b/rust/tests/tirx/index_map.rs
@@ -19,9 +19,9 @@
//! Tests of `tirx::index_map`.
-use tvm_tirx::ir::{dtype, PrimExpr, PrimType, Var};
+use tvm_tirx::ir::{dtype, PrimType};
use tvm_tirx::tirx::IndexMap;
-use tvm_tirx::tvm_ffi::{Array, DLDevice, DLDeviceType};
+use tvm_tirx::tvm_ffi::{DLDevice, DLDeviceType, ObjectRefCore};
use super::{as_i64, int32, op, prim_var, range};
use crate::common::{serial, structural_equal};
@@ -30,11 +30,10 @@ use crate::common::{serial, structural_equal};
fn index_maps() {
let _serial = serial();
let (i, j) = (prim_var("i", PrimType::int(32)), prim_var("j",
PrimType::int(32)));
- let transpose = IndexMap::new(
- Array::from_vec(vec![Var::from(&i), Var::from(&j)]),
- Array::from_vec(vec![PrimExpr::from(&j), PrimExpr::from(&i)]),
- None,
- );
+ let transpose = IndexMap::new([&i, &j], [&j, &i], None).unwrap();
+ assert!(transpose.inverse_index_map.is_none());
+ let involution = IndexMap::new([&i, &j], [&j, &i],
Some(&transpose)).unwrap();
+
assert!(involution.inverse_index_map.as_ref().unwrap().same_as(&transpose));
let mapped = transpose.map_indices([int32(1), int32(2)], None).unwrap();
assert_eq!((as_i64(&mapped.get(0).unwrap()),
as_i64(&mapped.get(1).unwrap())), (2, 1));
let shape = transpose.map_shape([int32(4), int32(8)], None).unwrap();
@@ -44,13 +43,11 @@ fn index_maps() {
let k = prim_var("k", PrimType::int(32));
let split = IndexMap::new(
- Array::from_vec(vec![Var::from(&k)]),
- Array::from_vec(vec![
- op("prim._OpFloorDiv", &k, int32(4)),
- op("prim._OpFloorMod", &k, int32(4)),
- ]),
+ [&k],
+ [op("prim._OpFloorDiv", &k, int32(4)), op("prim._OpFloorMod", &k,
int32(4))],
None,
- );
+ )
+ .unwrap();
let (unsplit, predicate) = split.non_surjective_inverse([range(14)],
None).unwrap();
assert_eq!(unsplit.final_indices.len(), 1);
assert!(!structural_equal(&predicate, &int32(0)));
diff --git a/rust/tests/tirx/layout.rs b/rust/tests/tirx/layout.rs
index bec4199ae..aea687855 100644
--- a/rust/tests/tirx/layout.rs
+++ b/rust/tests/tirx/layout.rs
@@ -19,7 +19,7 @@
//! Tests of `tirx::layout`: layout axes, tile layouts and swizzled layouts.
-use tvm_tirx::ir::Expr;
+use tvm_tirx::ir::PrimExpr;
use tvm_tirx::sym::Analyzer;
use tvm_tirx::tirx::{Axis, ComposeLayout, Iter, Layout, TileLayout};
use tvm_tirx::tvm_ffi::{Map, String as FfiString};
@@ -30,7 +30,7 @@ use crate::common::{global, serial};
#[test]
fn memory_axes() {
let _serial = serial();
- let m = Axis::get("m").unwrap();
+ let m = Axis::get("m");
assert!(m.is_memory().unwrap() && !m.is_thread().unwrap());
assert!(m.backend().unwrap().is_none());
assert!(m.get_scope().unwrap().is_none());
@@ -46,7 +46,7 @@ fn layouts() {
assert_eq!(as_i64(&layout.size(None).unwrap()), 64);
assert_eq!(as_i64(&layout.span(Some("m")).unwrap()), 64);
let key = FfiString::from("m");
- let at = |coord: Map<FfiString, Expr>|
as_i64(&coord.get(&key).unwrap().unwrap());
+ let at = |coord: Map<FfiString, PrimExpr>|
as_i64(&coord.get(&key).unwrap().unwrap());
assert_eq!(at(layout.apply([int32(1), int32(2)], &analyzer).unwrap()), 10);
assert_eq!(at(layout.apply_linear(int32(10), &analyzer).unwrap()), 10);
assert_eq!(at(layout.apply_with_shape([int32(10)], [int32(64)],
&analyzer).unwrap()), 10);
@@ -73,7 +73,7 @@ fn layouts() {
fn thread_layout_scope() {
let _serial = serial();
global("tirx.cuda.RegisterLayoutAxes").call_tuple(()).unwrap();
- let laneid = Axis::get("laneid").unwrap();
+ let laneid = Axis::get("laneid");
let lanes =
TileLayout::new([Iter::new(int32(32), int32(1), laneid)],
Vec::<Iter>::new(), &Map::new());
let (inner, outer) = lanes.get_scope().unwrap().unwrap();
diff --git a/rust/tests/tirx/main.rs b/rust/tests/tirx/main.rs
index a827743aa..5d5a43546 100644
--- a/rust/tests/tirx/main.rs
+++ b/rust/tests/tirx/main.rs
@@ -40,7 +40,7 @@ mod transform;
use std::sync::atomic::{AtomicUsize, Ordering};
use tvm_tirx::ir::{
- Location, UNKNOWN_LOC, BaseFunc, DictAttrs, Expr, GlobalInfo, GlobalVar,
IRModule, IntImm,
+ UnknownLoc, UNKNOWN_LOC, BaseFunc, DictAttrs, Expr, GlobalInfo, GlobalVar,
IRModule, IntImm,
PrimExpr, PrimType, PrimVar, Range, SeqStmt, Stmt, Type, Var,
};
use tvm_tirx::runtime::MemoryScope;
@@ -55,7 +55,7 @@ fn int32(value: i64) -> PrimExpr {
fn var(name: &str, ty: impl Into<Type>) -> Var {
global("ir.Var")
- .call_tuple((FfiString::from(name), Some(ty.into()),
Location::unknown()))
+ .call_tuple((FfiString::from(name), Some(ty.into()),
UnknownLoc::new()))
.unwrap()
.try_into()
.unwrap()
@@ -66,12 +66,12 @@ fn prim_var(name: &str, ty: impl Into<PrimType>) -> PrimVar
{
}
fn op(name: &str, a: impl Into<Expr>, b: impl Into<Expr>) -> PrimExpr {
- global(name).call_tuple((a.into(), b.into(),
Location::unknown())).unwrap().try_into().unwrap()
+ global(name).call_tuple((a.into(), b.into(),
UnknownLoc::new())).unwrap().try_into().unwrap()
}
fn range(extent: i64) -> Range {
global("ir.Range")
- .call_tuple((int32(0), Some(int32(extent)), Location::unknown()))
+ .call_tuple((int32(0), Some(int32(extent)), UnknownLoc::new()))
.unwrap()
.try_into()
.unwrap()
@@ -90,7 +90,7 @@ fn memory_scope() -> MemoryScope {
/// The row-major layout of an 8 x 8 tile.
fn tile_8x8() -> TileLayout {
- let m = Axis::get("m").unwrap();
+ let m = Axis::get("m");
TileLayout::new(
[Iter::new(int32(8), int32(8), m.clone()), Iter::new(int32(8),
int32(1), m)],
Vec::<Iter>::new(),
@@ -99,11 +99,11 @@ fn tile_8x8() -> TileLayout {
}
fn void_function(params: Vec<Var>) -> Function {
- let zero = global("ir.Evaluate").call_tuple((int32(0),
Location::unknown())).unwrap();
+ let zero = global("ir.Evaluate").call_tuple((int32(0),
UnknownLoc::new())).unwrap();
let body: SeqStmt = global("ir.SeqStmt")
.call_tuple((
Array::<Stmt>::from_vec(vec![zero.try_into().unwrap()]),
- Some(Location::unknown()),
+ Some(UnknownLoc::new()),
))
.unwrap()
.try_into()
@@ -114,10 +114,15 @@ fn void_function(params: Vec<Var>) -> Function {
fn module_of(name: &str, func: &Function) -> IRModule {
let gv: GlobalVar = global("ir.GlobalVar")
- .call_tuple((FfiString::from(name), Location::unknown()))
+ .call_tuple((FfiString::from(name), UnknownLoc::new()))
.unwrap()
.try_into()
.unwrap();
let functions: Map<GlobalVar, BaseFunc> = [(gv,
BaseFunc::from(func))].into_iter().collect();
- IRModule::from_functions(&functions, None, &Map::<FfiString,
Array<GlobalInfo>>::new()).unwrap()
+ IRModule::new(
+ &functions,
+ DictAttrs::new(Map::new()),
+ &Map::<FfiString, Array<GlobalInfo>>::new(),
+ )
+ .unwrap()
}
diff --git a/rust/tests/tirx/op.rs b/rust/tests/tirx/op.rs
index 7a3cb1c58..895f592e1 100644
--- a/rust/tests/tirx/op.rs
+++ b/rust/tests/tirx/op.rs
@@ -21,7 +21,7 @@
//! against the native functions they follow.
use tvm_tirx::ir::{
- Call, Expr, FloatImm, FuncType, IntImm, Location, Op, PrimExpr, PrimType,
PrimVar, PtrType,
+ Call, Expr, FloatImm, FuncType, IntImm, UnknownLoc, Op, PrimExpr,
PrimType, PrimVar, PtrType,
StringImm, Type, UNKNOWN_LOC,
};
use tvm_tirx::tirx::{fast_erf_float_expr, gpu_thread_return, logaddexp,
reinterpret};
@@ -80,7 +80,7 @@ fn logaddexp_agrees_with_cpp() {
];
for (x, y) in &pairs {
let native =
- global("tirx._OpLogAddExp").call_tuple((x.clone(), y.clone(),
Location::unknown()));
+ global("tirx._OpLogAddExp").call_tuple((x.clone(), y.clone(),
UnknownLoc::new()));
let what = format!("logaddexp({}, {})", repr(x), repr(y));
agree(&what, logaddexp(x, y), native);
}
@@ -108,7 +108,7 @@ fn reinterpret_agrees_with_cpp() {
];
for (ty, x) in cases {
let native =
- global("tirx.reinterpret").call_tuple((ty.clone(), x.clone(),
Location::unknown()));
+ global("tirx.reinterpret").call_tuple((ty.clone(), x.clone(),
UnknownLoc::new()));
let what = format!("reinterpret({}, {})", repr(&ty), repr(&x));
agree(&what, reinterpret(ty, x), native);
}
diff --git a/rust/tests/tirx/tensor.rs b/rust/tests/tirx/tensor.rs
index 0369eae99..69e501008 100644
--- a/rust/tests/tirx/tensor.rs
+++ b/rust/tests/tirx/tensor.rs
@@ -19,7 +19,7 @@
//! Tests of `tirx::tensor`: tensor types.
-use tvm_tirx::ir::{UNKNOWN_LOC, Location, Expr, PrimExpr, PrimType, StringImm};
+use tvm_tirx::ir::{UNKNOWN_LOC, UnknownLoc, Expr, PrimExpr, PrimType,
StringImm};
use tvm_tirx::tirx::{tensor_is_scalar, Layout, TensorType};
use tvm_tirx::tvm_ffi::Array;
@@ -71,7 +71,7 @@ fn tensor_types_agree_with_cpp() {
(&scope).into(),
(&alignment).into(),
(&layout).into(),
- (&Location::unknown()).into(),
+ (&UnknownLoc::new()).into(),
],
);
let rust =
@@ -88,7 +88,7 @@ fn tensor_types_agree_with_cpp() {
(&scope).into(),
(&alignment).into(),
(&layout).into(),
- (&Location::unknown()).into(),
+ (&UnknownLoc::new()).into(),
],
);
// A non-primitive extent is not a `PrimExpr`: the check comes before
`TensorType::new`.
diff --git a/rust/tests/tirx/tile/dispatch.rs b/rust/tests/tirx/tile/dispatch.rs
index f30f16130..0414c99c9 100644
--- a/rust/tests/tirx/tile/dispatch.rs
+++ b/rust/tests/tirx/tile/dispatch.rs
@@ -19,7 +19,7 @@
//! Tests of `tirx::tile::dispatch`: the context tile instructions are
dispatched in.
-use tvm_tirx::ir::{Location, UNKNOWN_LOC, PrimExpr, PrimType, Stmt, Var};
+use tvm_tirx::ir::{UnknownLoc, UNKNOWN_LOC, PrimExpr, PrimType, Stmt, Var};
use tvm_tirx::target::Target;
use tvm_tirx::tirx::tile::DispatchContext;
use tvm_tirx::tirx::{TensorType, ThreadScope};
@@ -56,7 +56,7 @@ fn dispatch_contexts() {
assert!(structural_equal(&cached.downcast::<Var>().unwrap(), &tx_var));
let stmt: Stmt = global("ir.Evaluate")
- .call_tuple((int32(0), Location::unknown()))
+ .call_tuple((int32(0), UnknownLoc::new()))
.unwrap()
.try_into()
.unwrap();
diff --git a/rust/tests/tirx/tile/instruction.rs
b/rust/tests/tirx/tile/instruction.rs
index 9e9de4b1b..0ef7d09ae 100644
--- a/rust/tests/tirx/tile/instruction.rs
+++ b/rust/tests/tirx/tile/instruction.rs
@@ -19,7 +19,7 @@
//! Tests of `tirx::tile::instruction`: the tensor regions and validators of
tile instructions.
-use tvm_tirx::ir::{Location, UNKNOWN_LOC, Expr, Op, PrimType};
+use tvm_tirx::ir::{UnknownLoc, UNKNOWN_LOC, Expr, Op, PrimType};
use tvm_tirx::tirx::tile::{as_tensor_region, configure_tensor_instruction};
use tvm_tirx::tirx::TensorType;
use tvm_tirx::tvm_ffi::Array;
@@ -35,7 +35,7 @@ fn tensor_regions() {
TensorType::new([int32(4)], PrimType::float(32), &scope, 0, None,
UNKNOWN_LOC).unwrap();
let tensor = var("T", &ty);
let load: Expr = global("ir.TensorLoad")
- .call_tuple((Expr::from(&tensor), Array::from_vec(vec![int32(1)]),
Location::unknown()))
+ .call_tuple((Expr::from(&tensor), Array::from_vec(vec![int32(1)]),
UnknownLoc::new()))
.unwrap()
.try_into()
.unwrap();
