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 413e24a8b [FIX][RUST] Follow the Ramp integer check and load the LLVM
backend in tests (#86)
413e24a8b is described below
commit 413e24a8beb6e154e7f3d18412ea8344c71e35e9
Author: Linzhang Li <[email protected]>
AuthorDate: Mon Oct 26 12:28:34 2026 -0400
[FIX][RUST] Follow the Ramp integer check and load the LLVM backend in
tests (#86)
Two recent changes on `main` broke four tests of the Rust crate (`cargo
test --manifest-path
rust/Cargo.toml`), which CI does not run yet. This PR brings the Rust
side and the tests back in
line with them.
- #64 restricts `Ramp` to 32-bit and 64-bit integer bases and strides.
The Rust `Ramp::new`, which
follows the C++ constructor, still accepted any scalar, so
`ir prim::vector_expr::vector_nodes_agree_with_cpp` failed on a float
base.
- #72 moves the LLVM code generator into `libtvm_tirx_llvm_backend`,
which `import tvm_tirx` loads
and the Rust crate does not, so `codegen.llvm_*` and `target.build.llvm`
were missing in
`codegen::llvm::llvm_queries`, `codegen build_module` and `target
build_by_kind`.
---
rust/src/ir/prim/op.rs | 2 +-
rust/src/ir/prim/vector_expr.rs | 29 ++++++++++++++++++++++-------
rust/tests/codegen/llvm.rs | 3 +++
rust/tests/codegen/main.rs | 3 ++-
rust/tests/common/mod.rs | 24 ++++++++++++++++++++++--
rust/tests/ir/prim/op.rs | 7 +++++--
rust/tests/ir/prim/vector_expr.rs | 7 ++++++-
rust/tests/target/main.rs | 3 ++-
8 files changed, 63 insertions(+), 15 deletions(-)
diff --git a/rust/src/ir/prim/op.rs b/rust/src/ir/prim/op.rs
index 33f1a8ee8..0aa5f9da2 100644
--- a/rust/src/ir/prim/op.rs
+++ b/rust/src/ir/prim/op.rs
@@ -687,7 +687,7 @@ pub fn cast_at(
// Casting the base and the stride is valid only when every lane is
the same integer in
// both types: widening must not remove an overflow of the source, and
a floating-point
// cast need not be affine.
- let integer = matches!(to.code, support::INT | support::UINT);
+ let integer = matches!(to.code, support::INT | support::UINT) &&
matches!(to.bits, 32 | 64);
if let (Some(base), Some(stride)) =
(const_fold::int(&ramp.base_), const_fold::int(&ramp.stride))
{
diff --git a/rust/src/ir/prim/vector_expr.rs b/rust/src/ir/prim/vector_expr.rs
index 9c44b114e..90f52a5d5 100644
--- a/rust/src/ir/prim/vector_expr.rs
+++ b/rust/src/ir/prim/vector_expr.rs
@@ -266,14 +266,29 @@ impl Ramp {
loc: impl Into<Location>,
) -> Result<Self> {
let (base, stride, lanes) = (base.into_expr()?, stride.into_expr()?,
lanes.into_expr()?);
- let base_dtype = support::prim_dtype(&base)?;
- let stride_dtype = support::prim_dtype(&stride)?;
+ let base_ty = support::prim_type(&base)?;
+ let stride_ty = support::prim_type(&stride)?;
support::prim_dtype(&lanes)?;
- if !support::is_scalar(base_dtype) {
- return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR,
"base_ty.IsScalar()", ""));
- }
- if !support::is_scalar(stride_dtype) {
- return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR,
"stride_ty.IsScalar()", ""));
+ let (base_dtype, stride_dtype) = (base_ty.dtype, stride_ty.dtype);
+ // Only 32-bit and 64-bit integers step through the lanes of a ramp.
+ let index = |dtype: DLDataType| {
+ support::is_scalar(dtype)
+ && matches!(dtype.code, support::INT | support::UINT)
+ && matches!(dtype.bits, 32 | 64)
+ };
+ for (field, ty) in [("base", &base_ty), ("stride", &stride_ty)] {
+ if !index(ty.dtype) {
+ let condition = format!(
+ "{field}_ty.IsScalar() && \
+ {field}_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt) && \
+ ({field}_ty.bits() == 32 || {field}_ty.bits() == 64)"
+ );
+ let message = format!(
+ "Ramp {field} must be scalar int32, int64, uint32, or
uint64, but got {}",
+ support::repr(ty)
+ );
+ return Err(support::check_failed(tvm_ffi::TYPE_ERROR,
&condition, &message));
+ }
}
let stride = match stride_dtype == base_dtype {
true => stride,
diff --git a/rust/tests/codegen/llvm.rs b/rust/tests/codegen/llvm.rs
index e74843e6f..d2696ccb8 100644
--- a/rust/tests/codegen/llvm.rs
+++ b/rust/tests/codegen/llvm.rs
@@ -22,8 +22,11 @@
use tvm_tirx::codegen::{self, llvm};
use tvm_tirx::target::Target;
+use crate::common::load_llvm_backend;
+
#[test]
fn llvm_queries() {
+ load_llvm_backend();
assert!(llvm::version_major().unwrap() >= 10);
let sqrt = llvm::lookup_intrinsic_id("llvm.sqrt").unwrap();
assert!(sqrt > 0);
diff --git a/rust/tests/codegen/main.rs b/rust/tests/codegen/main.rs
index 86856c24e..7d5389ea5 100644
--- a/rust/tests/codegen/main.rs
+++ b/rust/tests/codegen/main.rs
@@ -28,10 +28,11 @@ mod llvm;
use tvm_tirx::codegen;
use tvm_tirx::target::Target;
-use common::empty_module;
+use common::{empty_module, load_llvm_backend};
#[test]
fn build_module() {
+ load_llvm_backend();
let module = empty_module();
codegen::build_module(&module, Target::new("llvm").unwrap()).unwrap();
codegen::build_module(&module, Target::new("c").unwrap()).unwrap();
diff --git a/rust/tests/common/mod.rs b/rust/tests/common/mod.rs
index d4ffe00a5..1a350b0f3 100644
--- a/rust/tests/common/mod.rs
+++ b/rust/tests/common/mod.rs
@@ -21,13 +21,14 @@
#![allow(dead_code)]
+use std::env::consts::{DLL_PREFIX, DLL_SUFFIX};
use std::mem::size_of;
-use std::sync::{Mutex, MutexGuard};
+use std::sync::{Mutex, MutexGuard, Once};
use tvm_tirx::ir::{BaseFunc, Expr, GlobalInfo, GlobalVar, IRModule,
MissingTypeObj, PrimTypeObj};
use tvm_tirx::tvm_ffi::tvm_ffi_sys::{TVMFFIGetTypeInfo, TVMFFITypeInfo};
use tvm_tirx::tvm_ffi::{
- Any, AnyCompatible, AnyView, Array, Function, Map, ObjectCore,
ObjectRefCore, Result,
+ Any, AnyCompatible, AnyView, Array, Function, Map, Module, ObjectCore,
ObjectRefCore, Result,
String as FfiString,
};
@@ -39,6 +40,25 @@ pub fn serial() -> MutexGuard<'static, ()> {
LOCK.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
}
+/// Load the LLVM backend library, which registers the `codegen.llvm_*`
functions and the `llvm`
+/// target build, as `import tvm_tirx` does. It sits beside
`libtvm_tirx_compiler`, whose directory
+/// the build script puts on the loader path; a build with
`TVM_TIRX_LLVM_NATIVE_ONLY` names it
+/// `tvm_tirx_llvm_backend_native`.
+pub fn load_llvm_backend() {
+ static LOAD: Once = Once::new();
+ LOAD.call_once(|| {
+ let mut errors = Vec::new();
+ for name in ["tvm_tirx_llvm_backend", "tvm_tirx_llvm_backend_native"] {
+ match
Module::load_from_file(format!("{DLL_PREFIX}{name}{DLL_SUFFIX}")) {
+ // Unloading the library would unregister its functions.
+ Ok(backend) => return std::mem::forget(backend),
+ Err(err) => errors.push(err.to_string()),
+ }
+ }
+ panic!("could not load the LLVM backend: {}", errors.join("\n"));
+ });
+}
+
pub fn global(name: &str) -> Function {
Function::get_global(name).unwrap()
}
diff --git a/rust/tests/ir/prim/op.rs b/rust/tests/ir/prim/op.rs
index e7e4c0887..d19061a41 100644
--- a/rust/tests/ir/prim/op.rs
+++ b/rust/tests/ir/prim/op.rs
@@ -769,13 +769,15 @@ fn cast_folds_and_broadcasts_as_cpp() {
(Expr::from(&x), "int32x4"),
(Expr::from(&f), "int32x4"),
(imm("int32", 3), "int64x4"),
- // A broadcast is cast element by element, and so is a ramp of
constants cast to an
- // integer type when every lane fits both types; any other ramp
becomes a `Cast` node.
+ // A broadcast is cast element by element, and so is a ramp of
constants cast to a 32-bit
+ // or 64-bit integer type when every lane fits both types; any other
ramp becomes a `Cast`
+ // node.
(Broadcast::new(&f, 4, UNKNOWN_LOC).unwrap().into(), "int32x4"),
(constant.clone(), "int64x4"),
(overflowing.clone(), "int64x4"),
(ramp(x.clone().into()), "int64x4"),
(constant.clone(), "float32x4"),
+ (constant.clone(), "int16x4"),
];
for (value, name) in &cases {
assert!(same(value.cast(prim(name)), &native_cast(value, name)),
"{name}");
@@ -790,6 +792,7 @@ fn cast_folds_and_broadcasts_as_cpp() {
assert!(node(&overflowing, "int64x4").as_node::<CastObj>().is_some());
assert!(node(&ramp(x.clone().into()),
"int64x4").as_node::<CastObj>().is_some());
assert!(node(&constant, "float32x4").as_node::<CastObj>().is_some());
+ assert!(node(&constant, "int16x4").as_node::<CastObj>().is_some());
let raw: Expr = Cast::new(PrimType::int(64), &x,
UNKNOWN_LOC).unwrap().into();
assert!(same(x.cast(PrimType::int(64)), &raw));
diff --git a/rust/tests/ir/prim/vector_expr.rs
b/rust/tests/ir/prim/vector_expr.rs
index 39c7ff452..8ad02a28a 100644
--- a/rust/tests/ir/prim/vector_expr.rs
+++ b/rust/tests/ir/prim/vector_expr.rs
@@ -55,6 +55,8 @@ fn vector_nodes_agree_with_cpp() {
let _serial = serial();
let x = Expr::from(var("x", PrimType::int(32)));
let x64 = Expr::from(var("x64", PrimType::int(64)));
+ let x16 = Expr::from(var("x16", PrimType::int(16)));
+ let u32 = Expr::from(var("u32", PrimType::uint(32)));
let v4 = Expr::from(var("v", PrimType::int(32).with_lanes(4).unwrap()));
let bf = Expr::from(var("bf", PrimType::bfloat(32)));
let f = Expr::from(var("f", PrimType::float(32)));
@@ -82,7 +84,10 @@ fn vector_nodes_agree_with_cpp() {
("Ramp", x.clone(), Expr::from(1), lanes_4.clone()),
("Ramp(int64 lanes)", x.clone(), Expr::from(1),
IntImm::int64(4).into()),
("Ramp(int64 stride constant)", x.clone(), IntImm::int64(2).into(),
lanes_4.clone()),
- ("Ramp(float stride constant)", f.clone(), Expr::from(2),
lanes_4.clone()),
+ ("Ramp(uint32 base)", u32.clone(), Expr::from(1), lanes_4.clone()),
+ ("Ramp(int16 base)", x16.clone(), Expr::from(1), lanes_4.clone()),
+ ("Ramp(int16 stride)", x.clone(), x16.clone(), lanes_4.clone()),
+ ("Ramp(float base)", f.clone(), Expr::from(2), lanes_4.clone()),
("Ramp(int32 to int64 stride constant)", x64.clone(), Expr::from(-3),
lanes_4.clone()),
("Ramp(int64 stride)", x.clone(), x64.clone(), lanes_4.clone()),
("Ramp(int stride of float base)", f.clone(), x.clone(),
lanes_4.clone()),
diff --git a/rust/tests/target/main.rs b/rust/tests/target/main.rs
index d1ca39ec6..d7c62763f 100644
--- a/rust/tests/target/main.rs
+++ b/rust/tests/target/main.rs
@@ -26,7 +26,7 @@ mod common;
use tvm_tirx::target::{self, Target, TargetKind, TargetKindObj, TargetObj};
use tvm_tirx::tvm_ffi::{Any, Map, Object, ObjectRefCast, String as FfiString,
TypeIndex};
-use common::{assert_type_contracts, empty_module, global};
+use common::{assert_type_contracts, empty_module, global, load_llvm_backend};
fn config(kind: &str) -> Map<FfiString, Any> {
Map::from_iter([(FfiString::from("kind"),
Any::from(FfiString::from(kind)))])
@@ -122,6 +122,7 @@ fn tags() {
#[test]
fn build_by_kind() {
+ load_llvm_backend();
let c = Target::new("c").unwrap();
target::build("c", empty_module(), &c).unwrap();
target::build("llvm", empty_module(),
Target::new("llvm").unwrap()).unwrap();
