chore: import upstream snapshot with attribution

tests/conda/condarc

@@ -0,0 +1,42 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# See https://docs.conda.io/projects/conda/en/latest/configuration.html for details
+
+# remote_connect_timeout_secs (float)
+#   The number seconds conda will wait for your client to establish a
+#   connection to a remote url resource.
+#
+remote_connect_timeout_secs: 10
+
+# remote_max_retries (int)
+#   The maximum number of retries each HTTP connection should attempt.
+#
+remote_max_retries: 6
+
+# remote_backoff_factor (int)
+#   The factor determines the time HTTP connection should wait for
+#   attempt.
+#
+remote_backoff_factor: 5
+
+# remote_read_timeout_secs (float)
+#   Once conda has connected to a remote resource and sent an HTTP
+#   request, the read timeout is the number of seconds conda will wait for
+#   the server to send a response.
+#
+remote_read_timeout_secs: 60.0

tests/cpp/.gitignore

@@ -0,0 +1,3 @@
+unittest
+*.d
+*_test

tests/cpp/README.md

@@ -0,0 +1,23 @@
+<!--- Licensed to the Apache Software Foundation (ASF) under one -->
+<!--- or more contributor license agreements.  See the NOTICE file -->
+<!--- distributed with this work for additional information -->
+<!--- regarding copyright ownership.  The ASF licenses this file -->
+<!--- to you under the Apache License, Version 2.0 (the -->
+<!--- "License"); you may not use this file except in compliance -->
+<!--- with the License.  You may obtain a copy of the License at -->
+
+<!---   http://www.apache.org/licenses/LICENSE-2.0 -->
+
+<!--- Unless required by applicable law or agreed to in writing, -->
+<!--- software distributed under the License is distributed on an -->
+<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
+<!--- KIND, either express or implied.  See the License for the -->
+<!--- specific language governing permissions and limitations -->
+<!--- under the License. -->
+# tests/cpp
+
+This folder contains some unit tests for C++ utilities in the codebase.
+
+In principle we aim to do most compiler related tests in the python to
+bring more development velocity, and only use this folder for low-level unit-tests.
+All tests should finish fast and not dependent on a presence of an accelerator device.

tests/cpp/arith_simplify_test.cc

@@ -0,0 +1,117 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/arith/analyzer.h>
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/te/operation.h>
+
+TEST(Simplify, MinMax) {
+  tvm::arith::Analyzer ana;
+  auto x = tvm::te::var("x");
+  auto e1 = (tvm::max(x, 1) - tvm::max(x, 1));
+  auto e1s = ana->canonical_simplify(e1);
+  TVM_FFI_ICHECK(tvm::tirx::is_zero(e1s));
+
+  auto e2 = (x * tvm::min(x, 1)) - (x * tvm::min(x, 1));
+  auto e2s = ana->canonical_simplify(e2);
+  TVM_FFI_ICHECK(tvm::tirx::is_zero(e2s));
+}
+
+TEST(Simplify, Mul) {
+  tvm::arith::Analyzer ana;
+  auto x = tvm::te::var("x");
+  auto e = (x * x) - (x * x);
+  auto es = ana->canonical_simplify(e);
+  TVM_FFI_ICHECK(tvm::tirx::is_zero(es));
+}
+
+TEST(Simplify, Mod) {
+  tvm::arith::Analyzer ana;
+  auto x = tvm::IntImm::Int32(10);
+  auto y = tvm::IntImm::Int32(12);
+  // Mod::make is used instead of % to avoid constant folding during
+  // calling operator%(x,y). Mod::make doesn't try constant folding,
+  // and therefore, the constant folding will be attempted in CanonicalSimplify
+  auto mod = ana->canonical_simplify(tvm::tirx::Mod(x, y));
+  auto es = ana->canonical_simplify(mod - x);
+  TVM_FFI_ICHECK(tvm::tirx::is_zero(es));
+}
+
+TEST(AnalyzerObjectRef, CopySharesMutableState) {
+  tvm::arith::Analyzer analyzer;
+  tvm::arith::Analyzer copy = analyzer;
+  auto x = tvm::te::var("x");
+
+  copy->Bind(x, tvm::Range::FromMinExtent(0, 8));
+
+  TVM_FFI_ICHECK(analyzer->CanProve(x < 8));
+}
+
+TEST(AnalyzerObjectRef, ConstHandleRefCanMutateAnalyzerState) {
+  tvm::arith::Analyzer analyzer;
+  const tvm::arith::Analyzer& analyzer_ref = analyzer;
+  auto x = tvm::te::var("x");
+
+  analyzer_ref->Bind(x, tvm::Range::FromMinExtent(0, 8));
+
+  TVM_FFI_ICHECK(analyzer->CanProve(x < 8));
+}
+
+TEST(AnalyzerObjectRef, CloneIsIndependent) {
+  tvm::arith::Analyzer analyzer;
+  auto x = tvm::te::var("x");
+  auto y = tvm::te::var("y");
+
+  analyzer->Bind(x, tvm::Range::FromMinExtent(0, 8));
+  analyzer->modular_set.Update(x, tvm::arith::ModularSet(4, 0));
+
+  tvm::arith::Analyzer clone = analyzer->Clone();
+  TVM_FFI_ICHECK(clone->CanProve(x < 8));
+  TVM_FFI_ICHECK(clone->modular_set(x)->coeff == 4);
+
+  clone->Bind(y, tvm::Range::FromMinExtent(0, 4));
+  clone->modular_set.Update(x, tvm::arith::ModularSet(8, 0), true);
+  TVM_FFI_ICHECK(clone->CanProve(y < 4));
+  TVM_FFI_ICHECK(!analyzer->CanProve(y < 4));
+  TVM_FFI_ICHECK(analyzer->CanProve(x < 8));
+  TVM_FFI_ICHECK(analyzer->modular_set(x)->coeff == 4);
+  TVM_FFI_ICHECK(clone->modular_set(x)->coeff == 8);
+}
+
+TEST(ConstantFold, Broadcast) {
+  tvm::ffi::StructuralEqual checker;
+  auto i32x4 = tvm::tirx::Broadcast(tvm::IntImm::Int32(10), 4);
+  auto i64x4 = tvm::cast(i32x4.ty().WithBits(64), i32x4);
+  auto i64x4_expected = tvm::tirx::Broadcast(tvm::IntImm::Int64(10), 4);
+  ASSERT_TRUE(checker(i64x4, i64x4_expected));
+}
+
+TEST(ConstantFold, Ramp) {
+  tvm::ffi::StructuralEqual checker;
+  auto i32x4 = tvm::tirx::Ramp(tvm::IntImm::Int32(10), tvm::IntImm::Int32(1), 4);
+  auto i64x4 = tvm::cast(i32x4.ty().WithBits(64), i32x4);
+  auto i64x4_expected = tvm::tirx::Ramp(tvm::IntImm::Int64(10), tvm::IntImm::Int64(1), 4);
+  ASSERT_TRUE(checker(i64x4, i64x4_expected));
+
+  auto f32x4 = tvm::cast(tvm::PrimType::Float(32, 4), i32x4);
+  auto f32x4_expected = tvm::tirx::Cast(tvm::PrimType::Float(32, 4), i32x4);
+  ASSERT_TRUE(checker(f32x4, f32x4_expected));
+}

tests/cpp/config_schema_test.cc

@@ -0,0 +1,323 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/string.h>
+#include <tvm/ir/config_schema.h>
+
+using namespace tvm;
+using namespace tvm::ir;
+namespace refl = tvm::ffi::reflection;
+
+// Basic int/bool/string option resolution
+TEST(ConfigSchema, BasicTypes) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("num_threads");
+  schema.def_option<bool>("verbose");
+  schema.def_option<ffi::String>("name");
+
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"num_threads", int64_t(8)},
+      {"verbose", true},
+      {"name", ffi::String("test")},
+  };
+
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.at("num_threads").cast<int64_t>(), 8);
+  ASSERT_EQ(resolved.at("verbose").cast<bool>(), true);
+  ASSERT_EQ(resolved.at("name").cast<ffi::String>(), "test");
+}
+
+// Type coercion: correct type passes
+TEST(ConfigSchema, TypeCoercionPasses) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("count");
+
+  ffi::Map<ffi::String, ffi::Any> config = {{"count", int64_t(42)}};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.at("count").cast<int64_t>(), 42);
+}
+
+// Type mismatch rejection (TypeError)
+TEST(ConfigSchema, TypeMismatchRejection) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("count");
+
+  // Pass a string where int64_t is expected
+  ffi::Map<ffi::String, ffi::Any> config = {{"count", ffi::String("not_a_number")}};
+  try {
+    schema.Resolve(config);
+    FAIL() << "Expected TypeError";
+  } catch (const ffi::Error& e) {
+    EXPECT_EQ(e.kind(), "TypeError");
+    EXPECT_NE(std::string(e.message()).find("count"), std::string::npos);
+  }
+}
+
+TEST(ConfigSchema, TypeMismatchBoolForInt) {
+  ConfigSchema schema;
+  schema.def_option<bool>("flag");
+
+  // Pass a string where bool is expected
+  ffi::Map<ffi::String, ffi::Any> config = {{"flag", ffi::String("true")}};
+  try {
+    schema.Resolve(config);
+    FAIL() << "Expected TypeError";
+  } catch (const ffi::Error& e) {
+    EXPECT_EQ(e.kind(), "TypeError");
+    EXPECT_NE(std::string(e.message()).find("flag"), std::string::npos);
+  }
+}
+
+// Default value materialization via refl::DefaultValue
+TEST(ConfigSchema, DefaultValueMaterialization) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("max_threads", refl::DefaultValue(int64_t(256)));
+  schema.def_option<bool>("debug", refl::DefaultValue(false));
+
+  // Empty config -> defaults are applied
+  ffi::Map<ffi::String, ffi::Any> config = {};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.at("max_threads").cast<int64_t>(), 256);
+  ASSERT_EQ(resolved.at("debug").cast<bool>(), false);
+}
+
+// Default value materialization via refl::DefaultValue
+TEST(ConfigSchema, DefaultValueViaTrait) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("max_threads", refl::DefaultValue(int64_t(128)));
+
+  ffi::Map<ffi::String, ffi::Any> config = {};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.at("max_threads").cast<int64_t>(), 128);
+}
+
+// Explicit value overrides default
+TEST(ConfigSchema, ExplicitOverridesDefault) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("max_threads", refl::DefaultValue(int64_t(256)));
+
+  ffi::Map<ffi::String, ffi::Any> config = {{"max_threads", int64_t(512)}};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.at("max_threads").cast<int64_t>(), 512);
+}
+
+// Missing non-required option stays absent
+TEST(ConfigSchema, MissingNonRequired) {
+  ConfigSchema schema;
+  schema.def_option<ffi::String>("optional_name");  // no default
+
+  ffi::Map<ffi::String, ffi::Any> config = {};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.find("optional_name"), resolved.end());
+}
+
+// Unknown key rejection (ValueError)
+TEST(ConfigSchema, UnknownKeyRejection) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("known_key");
+
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"known_key", int64_t(1)},
+      {"unknown_key", ffi::String("surprise")},
+  };
+  try {
+    schema.Resolve(config);
+    FAIL() << "Expected ValueError";
+  } catch (const ffi::Error& e) {
+    EXPECT_EQ(e.kind(), "ValueError");
+    EXPECT_NE(std::string(e.message()).find("unknown_key"), std::string::npos);
+  }
+}
+
+// Unknown key allowed when error_on_unknown is false
+TEST(ConfigSchema, UnknownKeyAllowed) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("known_key");
+  schema.set_error_on_unknown(false);
+
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"known_key", int64_t(1)},
+      {"unknown_key", ffi::String("allowed")},
+  };
+  auto resolved = schema.Resolve(config);
+  // Unknown keys are silently dropped (not validated, not forwarded)
+  ASSERT_EQ(resolved.at("known_key").cast<int64_t>(), 1);
+}
+
+// Canonicalizer runs as final step
+TEST(ConfigSchema, CanonicalizerRunsLast) {
+  ConfigSchema schema;
+  schema.def_option<ffi::String>("mcpu");
+
+  auto canonicalizer = [](ffi::Map<ffi::String, ffi::Any> config) {
+    ffi::String mcpu = config.at("mcpu").cast<ffi::String>();
+    config.Set("mcpu", ffi::String("canonical_") + mcpu);
+    // Canonicalizer can add new keys
+    config.Set("feature.is_fast", true);
+    return config;
+  };
+  schema.set_canonicalizer(canonicalizer);
+
+  ffi::Map<ffi::String, ffi::Any> config = {{"mcpu", ffi::String("arm")}};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.at("mcpu").cast<ffi::String>(), "canonical_arm");
+  ASSERT_EQ(resolved.at("feature.is_fast").cast<bool>(), true);
+}
+
+// Option declaration order is preserved in output
+TEST(ConfigSchema, DeclarationOrderPreserved) {
+  ConfigSchema schema;
+  schema.def_option<ffi::String>("zebra");
+  schema.def_option<ffi::String>("alpha");
+  schema.def_option<ffi::String>("middle");
+
+  auto& options = schema.ListOptions();
+  ASSERT_EQ(options.size(), 3);
+  ASSERT_EQ(std::string(options[0].key), "zebra");
+  ASSERT_EQ(std::string(options[1].key), "alpha");
+  ASSERT_EQ(std::string(options[2].key), "middle");
+}
+
+// Doc string trait accepted without affecting behavior
+TEST(ConfigSchema, DocStringTrait) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("count", "Number of items to process");
+
+  ffi::Map<ffi::String, ffi::Any> config = {{"count", int64_t(5)}};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.at("count").cast<int64_t>(), 5);
+}
+
+// Array option type
+TEST(ConfigSchema, ArrayOptionType) {
+  ConfigSchema schema;
+  schema.def_option<ffi::Array<ffi::String>>("keys");
+
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"keys", ffi::Array<ffi::String>{"cpu", "gpu"}},
+  };
+  auto resolved = schema.Resolve(config);
+  auto keys = resolved.at("keys").cast<ffi::Array<ffi::String>>();
+  ASSERT_EQ(keys.size(), 2);
+  ASSERT_EQ(keys[0], "cpu");
+  ASSERT_EQ(keys[1], "gpu");
+}
+
+// Map option type
+TEST(ConfigSchema, MapOptionType) {
+  ConfigSchema schema;
+  schema.def_option<ffi::Map<ffi::String, int64_t>>("params");
+
+  ffi::Map<ffi::String, int64_t> params = {{"a", 1}, {"b", 2}};
+  ffi::Map<ffi::String, ffi::Any> config = {{"params", params}};
+  auto resolved = schema.Resolve(config);
+  auto result = resolved.at("params").cast<ffi::Map<ffi::String, int64_t>>();
+  ASSERT_EQ(result.size(), 2);
+  ASSERT_EQ(result["a"], 1);
+  ASSERT_EQ(result["b"], 2);
+}
+
+// Duplicate option key rejected
+TEST(ConfigSchema, DuplicateKeyRejected) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("key1");
+  try {
+    schema.def_option<int64_t>("key1");
+    FAIL() << "Expected ValueError";
+  } catch (const ffi::Error& e) {
+    EXPECT_EQ(e.kind(), "ValueError");
+    EXPECT_NE(std::string(e.message()).find("key1"), std::string::npos);
+  }
+}
+
+// ListOptions returns option entries in declaration order
+TEST(ConfigSchema, ListOptionsNames) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("no_default");
+  schema.def_option<bool>("with_default", refl::DefaultValue(true));
+
+  auto& options = schema.ListOptions();
+  ASSERT_EQ(options.size(), 2);
+  ASSERT_EQ(std::string(options[0].key), "no_default");
+  ASSERT_EQ(std::string(options[1].key), "with_default");
+  ASSERT_EQ(std::string(options[0].type_str), "int");
+  ASSERT_EQ(std::string(options[1].type_str), "bool");
+  // Check default_value
+  ASSERT_FALSE(options[0].has_default);
+  ASSERT_TRUE(options[1].has_default);
+  ASSERT_EQ(options[1].default_value.cast<bool>(), true);
+}
+
+// OptionEntry stores correct type strings
+TEST(ConfigSchema, OptionEntryTypeStr) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("count");
+  schema.def_option<bool>("flag");
+
+  auto& options = schema.ListOptions();
+  ASSERT_EQ(std::string(options[0].type_str), "int");
+  ASSERT_EQ(std::string(options[1].type_str), "bool");
+}
+
+// HasOption query
+TEST(ConfigSchema, HasOption) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("exists");
+  ASSERT_TRUE(schema.HasOption("exists"));
+  ASSERT_FALSE(schema.HasOption("does_not_exist"));
+}
+
+// Empty config with no options
+TEST(ConfigSchema, EmptySchemaEmptyConfig) {
+  ConfigSchema schema;
+  ffi::Map<ffi::String, ffi::Any> config = {};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.size(), 0);
+}
+
+// Canonicalizer with defaults
+TEST(ConfigSchema, CanonicalizerWithDefaults) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("threads", refl::DefaultValue(int64_t(4)));
+
+  auto canonicalizer = [](ffi::Map<ffi::String, ffi::Any> config) {
+    int64_t threads = config.at("threads").cast<int64_t>();
+    config.Set("threads", threads * 2);
+    return config;
+  };
+  schema.set_canonicalizer(canonicalizer);
+
+  // Empty config: default 4 applied, then canonicalizer doubles it
+  ffi::Map<ffi::String, ffi::Any> config = {};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.at("threads").cast<int64_t>(), 8);
+}
+
+// Combined doc + default traits
+TEST(ConfigSchema, DocAndDefaultTraits) {
+  ConfigSchema schema;
+  schema.def_option<int64_t>("num_cores", "CPU core count", refl::DefaultValue(int64_t(4)));
+
+  ffi::Map<ffi::String, ffi::Any> config = {};
+  auto resolved = schema.Resolve(config);
+  ASSERT_EQ(resolved.at("num_cores").cast<int64_t>(), 4);
+}

tests/cpp/expr_test.cc

@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/te/operation.h>
+
+TEST(Expr, Basic) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  PrimVar x("x");
+  auto z = max(x + 1 + 2, 100);
+  ffi::ObjectRef tmp = z;
+  PrimExpr zz = tmp.as_or_throw<PrimExpr>();
+  std::ostringstream os;
+  os << z;
+  TVM_FFI_ICHECK(zz.same_as(z));
+  TVM_FFI_ICHECK(os.str() == "T.max(x + 1 + 2, 100)");
+}
+
+TEST(Expr, VarTypeAnnotation) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  PrimVar x("x", PrimType::Float(32));
+  PrimVar y("y", PrimType::Float(32));
+  tvm::ffi::StructuralEqual checker;
+  TVM_FFI_ICHECK(checker(x.ty(), y.ty()));
+  TVM_FFI_ICHECK(checker(x->ty, y->ty));
+}
+
+TEST(Expr, PrimTypeBoolLanes) {
+  using namespace tvm;
+  PrimType boolx4 = PrimType::Bool(4);
+  TVM_FFI_ICHECK(boolx4.IsFixedLengthVector());
+  TVM_FFI_ICHECK(boolx4.MatchesCode(DLDataTypeCode::kDLBool));
+  TVM_FFI_ICHECK_EQ(boolx4.lanes(), 4);
+  TVM_FFI_ICHECK(boolx4.MatchesElementType(DLDataTypeCode::kDLBool, 8));
+}
+
+TEST(ExprNodeRef, Basic) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  PrimVar x("x");
+  PrimExpr z = max(x + 1 + 2, 100);
+  const tirx::MaxNode* op = z.as<tirx::MaxNode>();
+  TVM_FFI_ICHECK(ffi::GetRef<ffi::ObjectRef>(op).same_as(z));
+}

tests/cpp/ir_functor_test.cc

@@ -0,0 +1,379 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/ir/module.h>
+#include <tvm/ir/node_functor.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/builtin.h>
+#include <tvm/tirx/expr.h>
+#include <tvm/tirx/expr_functor.h>
+#include <tvm/tirx/function.h>
+#include <tvm/tirx/op.h>
+#include <tvm/tirx/stmt_functor.h>
+
+TEST(IRF, Basic) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  PrimVar x("x");
+  auto z = x + 1;
+
+  NodeFunctor<int(const ffi::ObjectRef& n, int b)> f;
+  f.set_dispatch<VarNode>([](const ffi::ObjectRef& n, int b) { return b; });
+  f.set_dispatch<AddNode>([](const ffi::ObjectRef& n, int b) { return b + 2; });
+  TVM_FFI_ICHECK_EQ(f(x, 2), 2);
+  TVM_FFI_ICHECK_EQ(f(z, 2), 4);
+}
+
+TEST(IRF, CountVar) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  int n_var = 0;
+  PrimVar x("x"), y("y");
+
+  auto z = x + 1 + y + y;
+  tirx::PostOrderVisit(z, [&n_var](const ffi::ObjectRef& n) {
+    if (n.as<VarNode>()) ++n_var;
+  });
+  TVM_FFI_ICHECK_EQ(n_var, 2);
+}
+
+TEST(IRF, PreOrderVisit) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  Stmt init =
+      IfThenElse(IntImm::Bool(true), Evaluate(IntImm::Int32(0)), Evaluate(IntImm::Int32(0)));
+  Stmt body = Evaluate(IntImm::Int32(1));
+  SBlock block(/*iter_vars=*/{}, /*reads=*/{},
+               /*writes=*/{}, /*name_hint=*/"block", /*body=*/body,
+               /*init=*/init);
+  bool init_visited = false;
+  bool stopped_at_if = true;
+  bool body_visited = false;
+  PreOrderVisit(block, [&](const ffi::ObjectRef& n) -> bool {
+    if (n->IsInstance<IfThenElseNode>()) {
+      init_visited = true;
+      return false;
+    }
+    if (const auto* eval = n.as<EvaluateNode>()) {
+      if (const auto* int_imm = eval->value.as<IntImmNode>()) {
+        if (int_imm->value == 0) {
+          stopped_at_if = false;
+        } else if (int_imm->value == 1) {
+          body_visited = true;
+        } else {
+          TVM_FFI_THROW(InternalError) << "Unreachable";
+        }
+      }
+    }
+    return true;
+  });
+  ASSERT_EQ(init_visited, true);
+  ASSERT_EQ(stopped_at_if, true);
+  ASSERT_EQ(body_visited, true);
+}
+
+TEST(IRF, ExprTransform) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  PrimVar x("x");
+  auto z = x + 1;
+
+  class MyExprFunctor : public tirx::ExprFunctor<int(const Expr&, int)> {
+   public:
+    int VisitExpr_(const VarNode* op, int b) final { return b; }
+    int VisitExpr_(const IntImmNode* op, int b) final { return op->value; }
+    int VisitExpr_(const AddNode* op, int b) final {
+      return VisitExpr(op->a, b) + VisitExpr(op->b, b);
+    }
+  };
+  MyExprFunctor f;
+  TVM_FFI_ICHECK_EQ(f(x, 2), 2);
+  TVM_FFI_ICHECK_EQ(f(z, 2), 3);
+  try {
+    f(z - 1, 2);
+    TVM_FFI_THROW(InternalError) << "should fail";
+  } catch (tvm::ffi::Error&) {
+  }
+}
+
+TEST(IRF, ExprVisit) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  PrimVar x("x");
+  auto z = x + 1;
+
+  class MyVisitor : public tirx::ExprFunctor<void(const Expr&)>,
+                    public tirx::StmtFunctor<void(const Stmt&)> {
+   public:
+    int count = 0;
+    // implementation
+    void VisitExpr_(const VarNode* op) final { ++count; }
+    void VisitExpr_(const IntImmNode* op) final {}
+    void VisitExpr_(const AddNode* op) final {
+      VisitExpr(op->a);
+      VisitExpr(op->b);
+    }
+    void VisitStmt_(const EvaluateNode* op) final { VisitExpr(op->value); }
+  };
+  MyVisitor v;
+  v.VisitStmt(Evaluate(z));
+  TVM_FFI_ICHECK_EQ(v.count, 1);
+}
+
+TEST(IRF, StmtVisitor) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  PrimVar x("x");
+  class MyVisitor : public StmtExprVisitor {
+   public:
+    int count = 0;
+    // implementation
+    void VisitExpr_(const VarNode* op) final { ++count; }
+  };
+  MyVisitor v;
+  auto fmaketest = [&]() {
+    auto z = x + 1;
+    Stmt eval_body = Evaluate(z);
+    PrimType dtype = PrimType::Float(32);
+    Var data_var("b", PointerType(dtype));
+    Buffer buf(data_var, dtype, {z, z}, {}, PrimExpr(), "b", 0, 0, BufferType::kDefault);
+    // AllocBuffer is flat (no body). Return as SeqStmt with eval.
+    return SeqStmt({AllocBuffer(buf), eval_body});
+  };
+  v(fmaketest());
+  // AllocBuffer visits buffer shape via VisitBufferDef.
+  // shape = {z, z} where z = x + 1, so x is visited twice from shape + once from eval = 3
+  TVM_FFI_ICHECK_EQ(v.count, 3);
+
+  {
+    // tests for block and block_realize
+    Stmt body = fmaketest();
+    PrimType dtype = PrimType::Float(32);
+    Var buf_var("b", PointerType(dtype));
+    Buffer buffer = decl_buffer({16});
+    body = SeqStmt({DeclBuffer(buffer), std::move(body)});
+    BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)});
+    MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region);
+
+    // construct block and block_realize
+    SBlock block = SBlock({}, {buffer_region}, {buffer_region}, "block", body, body, {},
+                          {match_buffer_region});
+    Stmt block_realize = SBlockRealize({}, IntImm::Bool(true), block);
+
+    v.count = 0;
+    v(block_realize);
+    // x visited in: reads range (1), writes range (1), match_buffers range (1),
+    // init DeclBuffer(0) + AllocBuffer shape(2) + Evaluate(1) = 3,
+    // body DeclBuffer(0) + AllocBuffer shape(2) + Evaluate(1) = 3.
+    // Total: 1 + 1 + 1 + 3 + 3 = 9.
+    TVM_FFI_ICHECK_EQ(v.count, 9);
+  }
+}
+
+TEST(IRF, StmtMutator) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  PrimVar x("x");
+
+  class MyVisitor : public tirx::StmtMutator, public tirx::ExprMutator {
+   public:
+    using StmtMutator::operator();
+    using ExprMutator::operator();
+
+   protected:
+    // implementation
+    Expr VisitExpr_(const AddNode* op) final { return op->a; }
+    Stmt VisitStmt_(const SeqStmtNode* op) final { return StmtMutator::VisitSeqStmt_(op, true); }
+    Expr VisitExpr(const Expr& expr) final { return ExprMutator::VisitExpr(expr); }
+  };
+  auto fmakealloc = [&]() {
+    auto z = x + 1;
+    PrimType dtype = PrimType::Float(32);
+    Var data_var("b", PointerType(dtype));
+    Buffer buf(data_var, dtype, {1, z}, {}, PrimExpr(), "b", 0, 0, BufferType::kDefault);
+    return AllocBuffer(buf);
+  };
+
+  auto fmakeif = [&]() {
+    auto z = x + 1;
+    Stmt body = Evaluate(z);
+    return IfThenElse(x, Evaluate(0), body);
+  };
+
+  MyVisitor v;
+  {
+    auto alloc = fmakealloc();
+    Stmt body2 = Evaluate(1);
+    auto* bufptr = alloc.as<AllocBufferNode>()->buffer.get();
+    ffi::Array<Stmt> arr{std::move(alloc), body2, body2};
+    auto* arrptr = arr.get();
+    arr.MutateByApply([&](Stmt s) { return v(std::move(s)); });
+    TVM_FFI_ICHECK(arr.get() == arrptr);
+    // buffer IS mutated now (AllocBuffer mutator visits buffer shape via VisitBufferDef)
+    // shape was {1, x+1}, mutator transforms x+1 -> x, so buffer changes
+    TVM_FFI_ICHECK(arr[0].as<AllocBufferNode>()->buffer.get() != bufptr);
+  }
+  {
+    ffi::Array<Stmt> arr{fmakealloc()};
+    // mutate array get reference by another one, trigger copy.
+    ffi::Array<Stmt> arr2 = arr;
+    auto* arrptr = arr.get();
+    arr.MutateByApply([&](Stmt s) { return v(std::move(s)); });
+    TVM_FFI_ICHECK(arr.get() != arrptr);
+    // buffer is mutated in arr but not in arr2
+    TVM_FFI_ICHECK(arr[0].as<AllocBufferNode>()->buffer.get() !=
+                   arr2[0].as<AllocBufferNode>()->buffer.get());
+    // mutate but no content change.
+    arr2 = arr;
+    arr.MutateByApply([&](Stmt s) { return v(std::move(s)); });
+    TVM_FFI_ICHECK(arr2.get() == arr.get());
+  }
+  {
+    ffi::Array<Stmt> arr{fmakeif()};
+    arr.MutateByApply([&](Stmt s) { return v(std::move(s)); });
+    TVM_FFI_ICHECK(arr[0].as<IfThenElseNode>()->else_case.as<EvaluateNode>()->value.same_as(x));
+    // mutate but no content change.
+    auto arr2 = arr;
+    arr.MutateByApply([&](Stmt s) { return v(std::move(s)); });
+    TVM_FFI_ICHECK(arr2.get() == arr.get());
+  }
+
+  {
+    auto body =
+        Evaluate(Call(PrimType::Int(32), builtin::call_extern(), {StringImm("xyz"), x + 1}));
+    auto res = v(std::move(body));
+    TVM_FFI_ICHECK(res.as<EvaluateNode>()->value.as<CallNode>()->args[1].same_as(x));
+  }
+  {
+    Stmt body = fmakealloc();
+    Stmt body2 = Evaluate(1);
+    auto* ref2 = body2.get();
+    auto* bufptr = body.as<AllocBufferNode>()->buffer.get();
+    // construct a recursive SeqStmt.
+    body = SeqStmt({body, body2});
+    body = SeqStmt({body, body2});
+    body = v(std::move(body));
+    // the seq get flattened
+    TVM_FFI_ICHECK(body.as<SeqStmtNode>()->size() == 3);
+    // buffer is now mutated (shape x+1 -> x via VisitBufferDef)
+    TVM_FFI_ICHECK(body.as<SeqStmtNode>()->seq[0].as<AllocBufferNode>()->buffer.get() != bufptr);
+    TVM_FFI_ICHECK(body.as<SeqStmtNode>()->seq[1].get() == ref2);
+  }
+
+  {
+    // Cannot cow because of bref
+    Stmt body = fmakealloc();
+    Stmt body2 = Evaluate(1);
+    // construct a recursive SeqStmt.
+    body = SeqStmt({body, body2});
+    auto bref = body;
+    body = SeqStmt({body, body2});
+    body = v(std::move(body));
+    // the seq get flattened
+    TVM_FFI_ICHECK(body.as<SeqStmtNode>()->size() == 3);
+    // buffer is mutated (shape x+1 -> x via VisitBufferDef)
+    TVM_FFI_ICHECK(body.as<SeqStmtNode>()->seq[0].as<AllocBufferNode>() != nullptr);
+    // bref still holds the old SeqStmt (not shared with new one due to copy)
+    TVM_FFI_ICHECK(!bref.same_as(body));
+  }
+
+  {
+    // tests for block and block_realize
+    // AllocBuffer and DeclBuffer are flat (no body), placed as siblings in SeqStmt
+    Stmt eval_body = Evaluate(x + 1);
+    Buffer buffer = decl_buffer({16});
+    Stmt decl = DeclBuffer(buffer);
+    Stmt alloc = fmakealloc();
+    // body is: DeclBuffer, AllocBuffer, Evaluate
+    Stmt body = SeqStmt({decl, alloc, eval_body});
+    BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)});
+    MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region);
+    // construct block and block_realize
+    SBlock block = SBlock({}, {buffer_region}, {buffer_region}, "block", body, body, {},
+                          {match_buffer_region});
+    Stmt block_realize = SBlockRealize({}, IntImm::Bool(true), block);
+    body = v(std::move(block_realize));
+    // the body should be changed
+    SBlock new_block = body.as<SBlockRealizeNode>()->block;
+    // body is a SeqStmt; the Evaluate(x+1) -> Evaluate(x)
+    auto* seq = new_block->body.as<SeqStmtNode>();
+    TVM_FFI_ICHECK(seq != nullptr);
+    TVM_FFI_ICHECK(seq->seq[2].as<EvaluateNode>()->value.same_as(x));
+    auto* init_seq = new_block->init.value().as<SeqStmtNode>();
+    TVM_FFI_ICHECK(init_seq != nullptr);
+    TVM_FFI_ICHECK(init_seq->seq[2].as<EvaluateNode>()->value.same_as(x));
+    // buffer region min is mutated: x+1 -> x
+    TVM_FFI_ICHECK(new_block->reads[0]->region[0]->min.same_as(x));
+    TVM_FFI_ICHECK(new_block->writes[0]->region[0]->min.same_as(x));
+    TVM_FFI_ICHECK(new_block->match_buffers[0]->source->region[0]->min.same_as(x));
+  }
+}
+
+TEST(IRF, Substitute) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  PrimType dtype = PrimType::Float(32);
+  Var x("x", PointerType(dtype, ""));
+  PrimVar n("n", PrimType::Int(32));
+
+  auto fmakebuffer = [&]() {
+    return Buffer{/*data=*/x,
+                  /*dtype=*/PrimType::Float(32),
+                  /*shape=*/{n},
+                  /*strides=*/{},
+                  /*elem_offset=*/PrimExpr(),
+                  /*name=*/"buf",
+                  /*data_alignment=*/1,
+                  /*offset_factor=*/1,
+                  /*buffer_type=*/BufferType::kDefault};
+  };
+
+  {
+    // test substitute buffer data var and shape var via DeclBuffer
+    Var y = x.CopyWithSuffix("subst");
+    PrimVar m("m", PrimType::Int(32));
+    Buffer buffer = fmakebuffer();
+    Stmt store = BufferStore(buffer, FloatImm(dtype, 0), {IntImm::Int32(0)});
+    Stmt decl = SeqStmt({DeclBuffer(buffer), store});
+    auto f_subst = [&](const Var& var) -> ffi::Optional<Expr> {
+      if (var.same_as(x)) return Expr(y);
+      if (var.same_as(n)) return Expr(m);
+      return std::nullopt;
+    };
+    Stmt new_decl = Substitute(decl, f_subst);
+    auto* seq_node = new_decl.as<SeqStmtNode>();
+    TVM_FFI_ICHECK(seq_node != nullptr);
+    auto* decl_node = seq_node->seq[0].as<DeclBufferNode>();
+    TVM_FFI_ICHECK(decl_node != nullptr);
+    TVM_FFI_ICHECK(decl_node->buffer->data.same_as(y));
+    TVM_FFI_ICHECK(decl_node->buffer->shape[0].same_as(m));
+  }
+
+  {
+    // test identity substitution on expression
+    Buffer buffer = fmakebuffer();
+    PrimExpr expr = BufferLoad(buffer, {IntImm::Int32(0)});
+    auto f_subst = [&](const Var& var) -> ffi::Optional<Expr> { return Expr(var); };
+    PrimExpr new_expr = Substitute(expr, f_subst);
+    // the expression is not changed
+    TVM_FFI_ICHECK(new_expr.same_as(expr));
+  }
+}

tests/cpp/ndarray_test.cc

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/runtime/tensor.h>
+
+using namespace tvm;
+
+TEST(TensorTest, IsContiguous_ContiguousStride) {
+  auto array = runtime::Tensor::Empty({5, 10}, DLDataType{kDLFloat, 32, 1}, {kDLCPU});
+  DLManagedTensor* managed_tensor = array.ToDLPack();
+
+  int64_t strides[] = {10, 1};
+  managed_tensor->dl_tensor.strides = strides;
+
+  TVM_FFI_ICHECK(ffi::IsContiguous(managed_tensor->dl_tensor));
+
+  managed_tensor->deleter(managed_tensor);
+}
+
+TEST(TensorTest, IsContiguous_NullStride) {
+  auto array = runtime::Tensor::Empty({5, 10}, DLDataType{kDLFloat, 32, 1}, {kDLCPU});
+  DLManagedTensor* managed_tensor = array.ToDLPack();
+
+  managed_tensor->dl_tensor.strides = nullptr;
+
+  TVM_FFI_ICHECK(ffi::IsContiguous(managed_tensor->dl_tensor));
+
+  managed_tensor->deleter(managed_tensor);
+}
+
+TEST(TensorTest, IsContiguous_AnyStrideForSingular) {
+  auto array = runtime::Tensor::Empty({5, 1, 10}, DLDataType{kDLFloat, 32, 1}, {kDLCPU});
+  DLManagedTensor* managed_tensor = array.ToDLPack();
+
+  int64_t strides[] = {10, 1, 1};  // strides[1] is normalized to 1 because shape[1] == 1.
+  managed_tensor->dl_tensor.strides = strides;
+
+  TVM_FFI_ICHECK(ffi::IsContiguous(managed_tensor->dl_tensor));
+
+  managed_tensor->dl_tensor.strides = nullptr;
+  managed_tensor->deleter(managed_tensor);
+}
+
+TEST(TensorTest, IsContiguous_UncontiguousStride) {
+  auto array = runtime::Tensor::Empty({5, 1, 10}, DLDataType{kDLFloat, 32, 1}, {kDLCPU});
+  DLManagedTensor* managed_tensor = array.ToDLPack();
+
+  int64_t strides[] = {1, 1, 1};
+  managed_tensor->dl_tensor.strides = strides;
+
+  TVM_FFI_ICHECK(!ffi::IsContiguous(managed_tensor->dl_tensor));
+
+  managed_tensor->dl_tensor.strides = nullptr;
+  managed_tensor->deleter(managed_tensor);
+}

tests/cpp/nested_msg_test.cc

@@ -0,0 +1,343 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/relax/block_builder.h>
+#include <tvm/relax/nested_msg.h>
+#include <tvm/relax/type.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/tirx/expr.h>
+
+#include <algorithm>
+#include <array>
+#include <cstring>
+#include <functional>
+#include <iterator>
+#include <new>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+using namespace tvm;
+using namespace tvm::relax;
+
+namespace {
+
+TensorType ScalarTensorType(PrimType dtype) {
+  auto n = tvm::ffi::make_object<TensorTypeNode>();
+  n->dtype = std::move(dtype);
+  n->ndim = 0;
+  return TensorType(n);
+}
+
+}  // namespace
+
+TEST(NestedMsg, Basic) {
+  // start with no annotation
+  relax::Var x("x", std::nullopt), y("y", std::nullopt);
+
+  // constructor from array, T and nullopt.
+  NestedMsg<relax::Expr> msg({x, std::nullopt, x});
+
+  EXPECT_TRUE(msg.IsNested());
+  EXPECT_FALSE(msg.IsLeaf());
+  EXPECT_TRUE(msg != nullptr);
+
+  EXPECT_ANY_THROW(msg.LeafValue());
+
+  auto arr = msg.NestedArray();
+  EXPECT_TRUE(arr[0].LeafValue().same_as(x));
+  EXPECT_TRUE(arr[1] == nullptr);
+  EXPECT_TRUE(arr[1].IsNull());
+
+  EXPECT_TRUE(arr[2].LeafValue().same_as(x));
+
+  auto a0 = arr[0];
+  EXPECT_TRUE(a0.IsLeaf());
+
+  // assignment
+  // assign null
+  a0 = std::nullopt;
+  EXPECT_TRUE(a0 == nullptr);
+
+  // assign array
+  a0 = {x, {x, std::nullopt, y}};
+  EXPECT_TRUE(a0.IsNested());
+  auto t0 = a0.NestedArray()[1];
+  EXPECT_TRUE(t0.IsNested());
+  EXPECT_TRUE(t0.NestedArray()[2].LeafValue().same_as(y));
+
+  // assign leaf
+  a0 = x;
+
+  EXPECT_TRUE(a0.IsLeaf());
+  EXPECT_TRUE(a0.LeafValue().same_as(x));
+}
+
+TEST(NestedMsg, IntAndAny) {
+  NestedMsg<int64_t> msg({1, std::nullopt, 2});
+  Any any_msg = msg;
+  NestedMsg<int64_t> msg2 = any_msg.cast<NestedMsg<int64_t>>();
+
+  EXPECT_TRUE(msg2.IsNested());
+  EXPECT_EQ(msg2.NestedArray()[0].LeafValue(), 1);
+  EXPECT_TRUE(msg2.NestedArray()[1].IsNull());
+  EXPECT_EQ(msg2.NestedArray()[2].LeafValue(), 2);
+}
+
+TEST(NestedMsg, ForEachLeaf) {
+  relax::Var x("x", std::nullopt), y("y", std::nullopt);
+  NestedMsg<Expr> msg = {x, {x, y}, std::nullopt, {x, {x, y}}};
+
+  int x_count = 0, y_count = 0;
+
+  ForEachLeaf(msg, [&](const Expr& v) {
+    if (v.same_as(x)) ++x_count;
+    if (v.same_as(y)) ++y_count;
+  });
+  EXPECT_EQ(x_count, 4);
+  EXPECT_EQ(y_count, 2);
+}
+
+TEST(NestedMsg, Equal) {
+  relax::Var x("x", std::nullopt), y("y", std::nullopt);
+  relax::Var z("z", std::nullopt);
+
+  auto fequal = [](Expr lhs, Expr rhs) { return lhs.same_as(rhs); };
+
+  using M = NestedMsg<relax::Expr>;
+
+  EXPECT_TRUE(Equal(M(std::nullopt), M(std::nullopt), fequal));
+
+  EXPECT_TRUE(Equal(M(x), M(x), fequal));
+
+  EXPECT_TRUE(Equal(M({x, y}), M({x, y}), fequal));
+
+  EXPECT_TRUE(Equal(M({x, std::nullopt}), M({x, std::nullopt}), fequal));
+
+  EXPECT_TRUE(Equal(M({x, {std::nullopt, y}}), M({x, {std::nullopt, y}}), fequal));
+
+  EXPECT_TRUE(Equal(M({x, {std::nullopt, y}, {x, z}}), M({x, {std::nullopt, y}, {x, z}}), fequal));
+
+  // type mismatch
+  EXPECT_FALSE(Equal(M({x, {std::nullopt, y}, x}), M({x, {std::nullopt, y}, {x, z}}), fequal));
+
+  EXPECT_FALSE(Equal(M({x, {std::nullopt, y}, {x, std::nullopt}}),
+                     M({x, {std::nullopt, y}, {x, z}}), fequal));
+
+  EXPECT_FALSE(Equal(M({x, {std::nullopt, y}}), M({x, {std::nullopt, y}, {x, z}}), fequal));
+
+  EXPECT_FALSE(Equal(M(x), M(std::nullopt), fequal));
+
+  EXPECT_FALSE(Equal(M(std::nullopt), M(x), fequal));
+
+  EXPECT_FALSE(Equal(M(x), M(ffi::Array<M>({x})), fequal));
+
+  EXPECT_FALSE(Equal(M(ffi::Array<M>({x})), M(x), fequal));
+}
+
+TEST(NestedMsg, MapAndDecompose) {
+  relax::Var x("x", PrimType::Int(16));
+  relax::Var y("y", PrimType::Int(32));
+  relax::Var z("z", PrimType::Int(64));
+
+  BlockBuilder bb = BlockBuilder::Create(std::nullopt);
+  relax::Expr t0 = bb->Normalize(Tuple({x, y}));
+  relax::Expr t1 = bb->Normalize(Tuple({t0, x, z, t0}));
+
+  auto c0 = IntImm::Int32(0);
+  auto c1 = IntImm::Int32(1);
+  auto c2 = IntImm::Int32(2);
+
+  auto output = MapToNestedMsg<IntImm>(t1, [&](Expr value) {
+    if (value.same_as(x)) return c0;
+    if (value.same_as(y)) return c1;
+    return c2;
+  });
+
+  NestedMsg<IntImm> expected = {{c0, c1}, c0, c2, {c0, c1}};
+
+  EXPECT_TRUE(Equal(output, expected,
+                    [](IntImm lhs, IntImm rhs) -> bool { return lhs->value == rhs->value; }));
+
+  auto output2 = MapToNestedMsg<IntImm>(GetType(t1), [&](Type ty) -> NestedMsg<IntImm> {
+    const auto* prim_ty = ty.as<PrimTypeNode>();
+    if (prim_ty == nullptr) return std::nullopt;
+    int bits = prim_ty->dtype.bits;
+    if (bits == 16) return c0;
+    if (bits == 32) return c1;
+    if (bits == 64) return c2;
+    return std::nullopt;
+  });
+
+  EXPECT_TRUE(Equal(output2, expected,
+                    [](IntImm lhs, IntImm rhs) -> bool { return lhs->value == rhs->value; }));
+
+  int x_count = 0, y_count = 0, z_count = 0;
+
+  DecomposeNestedMsg(t1, expected, [&](Expr value, NestedMsg<IntImm> msg) {
+    if (value.same_as(x)) {
+      EXPECT_TRUE(msg.LeafValue().same_as(c0));
+      ++x_count;
+    } else if (value.same_as(y)) {
+      EXPECT_TRUE(msg.LeafValue().same_as(c1));
+      ++y_count;
+    } else {
+      EXPECT_TRUE(msg.LeafValue().same_as(c2));
+      ++z_count;
+    }
+  });
+  EXPECT_EQ(x_count, 3);
+  EXPECT_EQ(y_count, 2);
+  EXPECT_EQ(z_count, 1);
+}
+
+TEST(NestedMsg, MapToNestedMsgByType) {
+  auto sf0 = ScalarTensorType(PrimType::Float(32));
+  auto sf1 = TupleType({sf0, sf0});
+  auto sf2 = TupleType({sf0, sf0});
+  auto x = relax::Var("x", TupleType({sf1, sf2, sf0}));
+
+  auto msg = MapToNestedMsgByType<Expr>(x, [](Expr value) { return value; });
+
+  EXPECT_TRUE(msg.IsNested());
+  auto arr = msg.NestedArray();
+
+  EXPECT_TRUE(arr[1].IsNested());
+  auto arr1 = arr[1].NestedArray();
+
+  EXPECT_TRUE(arr1[0].IsLeaf());
+  EXPECT_TRUE(
+      tvm::ffi::StructuralEqual()(arr1[0].LeafValue(), TupleGetItem(TupleGetItem(x, 1), 0)));
+
+  EXPECT_TRUE(arr[2].IsLeaf());
+  EXPECT_TRUE(tvm::ffi::StructuralEqual()(arr[2].LeafValue(), TupleGetItem(x, 2)));
+}
+
+TEST(NestedMsg, NestedMsgToExpr) {
+  auto sf0 = ScalarTensorType(PrimType::Float(32));
+  auto sf1 = TupleType({sf0, sf0});
+
+  auto c0 = IntImm::Int32(0);
+  auto c1 = IntImm::Int32(1);
+  auto c2 = IntImm::Int32(2);
+
+  relax::Var x("x", sf0), y("y", sf0), z("z", sf0);
+
+  NestedMsg<IntImm> msg = {c0, {c0, c1}, {c0, {c1, c2}}};
+  auto expr = NestedMsgToExpr<IntImm>(msg, [&](ffi::Optional<IntImm> leaf) {
+    TVM_FFI_ICHECK(leaf.has_value());
+    int value = leaf.value()->value;
+    switch (value) {
+      case 0:
+        return x;
+      case 1:
+        return y;
+      default:
+        return z;
+    }
+  });
+
+  Expr expected = Tuple({x, Tuple({x, y}), Tuple({x, Tuple({y, z})})});
+  EXPECT_TRUE(tvm::ffi::StructuralEqual()(expr, expected));
+
+  // test simplified
+  relax::Var t("t", sf1);
+  NestedMsg<Expr> msg1 = {TupleGetItem(t, 0), TupleGetItem(t, 1)};
+  auto expr1 = NestedMsgToExpr<Expr>(msg1, [](ffi::Optional<Expr> leaf) { return leaf.value(); });
+  EXPECT_TRUE(tvm::ffi::StructuralEqual()(expr1, t));
+}
+
+TEST(NestedMsg, CombineNestedMsg) {
+  auto c0 = IntImm::Int32(0);
+  auto c1 = IntImm::Int32(1);
+  auto c2 = IntImm::Int32(2);
+
+  NestedMsg<IntImm> lhs = {c0, {c0, c1}, std::nullopt, {c0, {c1, c2}}};
+  NestedMsg<IntImm> rhs = {c1, {c2, std::nullopt}, std::nullopt, {c1, {c2, c2}}};
+  NestedMsg<IntImm> expected = {c1, {c2, c1}, std::nullopt, {c1, {c2, c2}}};
+
+  auto output = CombineNestedMsg(lhs, rhs, [](IntImm x, IntImm y) {
+    if (x->value > y->value) return x;
+    return y;
+  });
+
+  EXPECT_TRUE(Equal(output, expected,
+                    [](IntImm lhs, IntImm rhs) -> bool { return lhs->value == rhs->value; }));
+}
+
+TEST(NestedMsg, MapNestedMsg) {
+  auto c0 = IntImm::Int32(0);
+  auto c1 = IntImm::Int32(1);
+  auto c2 = IntImm::Int32(2);
+  auto c3 = IntImm::Int32(3);
+
+  NestedMsg<IntImm> msg = {c0, {c0, c1}, std::nullopt, {c0, {c2, c1}}};
+  NestedMsg<IntImm> expected = {c3, {c3, std::nullopt}, std::nullopt, {c3, {c2, std::nullopt}}};
+
+  auto output = MapNestedMsg(msg, [](IntImm x) {
+    if (x->value == 0) {
+      return NestedMsg<IntImm>(IntImm::Int32(3));
+    } else if (x->value == 1) {
+      return NestedMsg<IntImm>();
+    } else {
+      return NestedMsg<IntImm>(x);
+    }
+  });
+
+  EXPECT_TRUE(Equal(output, expected,
+                    [](IntImm lhs, IntImm rhs) -> bool { return lhs->value == rhs->value; }));
+}
+
+TEST(NestedMsg, TransformTupleLeaf) {
+  auto c0 = IntImm::Int32(0);
+  auto c1 = IntImm::Int32(1);
+  auto c2 = IntImm::Int32(2);
+  using NInt = NestedMsg<IntImm>;
+
+  NInt msg1 = {c0, {c0, c1}, c2, {c0, {c1, c2}}};
+  NInt msg2 = {c1, {c2, c0}, c2, {c1, {c2, c0}}};
+
+  PrimType s = PrimType::Int(32);
+  relax::Var x("x", s), y("y", s), z("z", s);
+  BlockBuilder bb = BlockBuilder::Create(std::nullopt);
+  Expr expr = bb->Normalize(Tuple({x, Tuple({x, x}), x, Tuple({x, Tuple({x, x})})}));
+
+  auto ftransleaf = [&](Expr value, std::array<NInt, 2> msgs) -> Expr {
+    int lhs = msgs[0].LeafValue().as_or_throw<IntImm>()->value;
+    int rhs = msgs[1].LeafValue().as_or_throw<IntImm>()->value;
+    if (lhs > rhs)
+      return z;
+    else if (lhs == rhs)
+      return value;
+    else
+      return y;
+  };
+
+  Expr expected = Tuple({y, Tuple({y, z}), x, Tuple({y, Tuple({y, z})})});
+
+  EXPECT_TRUE(tvm::ffi::StructuralEqual()(
+      TransformTupleLeaf(expr, std::array<NInt, 2>({msg1, msg2}), ftransleaf), expected));
+
+  EXPECT_TRUE(
+      expr.same_as(TransformTupleLeaf(expr, std::array<NInt, 2>({msg1, msg1}), ftransleaf)));
+}

tests/cpp/object_protocol_test.cc

@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/ffi/memory.h>
+#include <tvm/runtime/logging.h>
+
+namespace tvm {
+namespace test {
+
+using namespace tvm::runtime;
+
+class ObjBase : public ffi::Object {
+ public:
+  // dynamically allocate slow
+  static constexpr const uint32_t _type_child_slots = 1;
+  TVM_FFI_DECLARE_OBJECT_INFO("test.ObjBase", ObjBase, ffi::Object);
+};
+
+class ObjA : public ObjBase {
+ public:
+  static constexpr const uint32_t _type_child_slots = 0;
+  TVM_FFI_DECLARE_OBJECT_INFO("test.ObjA", ObjA, ObjBase);
+};
+
+class ObjB : public ObjBase {
+ public:
+  static constexpr const uint32_t _type_child_slots = 0;
+  TVM_FFI_DECLARE_OBJECT_INFO("test.ObjB", ObjB, ObjBase);
+};
+
+class ObjAA : public ObjA {
+ public:
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.ObjAA", ObjAA, ObjA);
+};
+
+}  // namespace test
+}  // namespace tvm
+
+TEST(ObjectHierachy, Basic) {
+  using namespace tvm;
+  using namespace tvm::runtime;
+  using namespace tvm::test;
+  using namespace tvm::ffi;
+
+  ffi::ObjectRef refA(ffi::make_object<ObjA>());
+  TVM_FFI_ICHECK_EQ(refA->type_index(), ObjA::RuntimeTypeIndex());
+  TVM_FFI_ICHECK(refA.as<ffi::Object>() != nullptr);
+  TVM_FFI_ICHECK(refA.as<ObjA>() != nullptr);
+  TVM_FFI_ICHECK(refA.as<ObjBase>() != nullptr);
+  TVM_FFI_ICHECK(refA.as<ObjB>() == nullptr);
+  TVM_FFI_ICHECK(refA.as<ObjAA>() == nullptr);
+
+  ffi::ObjectRef refAA(ffi::make_object<ObjAA>());
+  TVM_FFI_ICHECK_EQ(refAA->type_index(), ObjAA::RuntimeTypeIndex());
+  TVM_FFI_ICHECK(refAA.as<ffi::Object>() != nullptr);
+  TVM_FFI_ICHECK(refAA.as<ObjBase>() != nullptr);
+  TVM_FFI_ICHECK(refAA.as<ObjA>() != nullptr);
+  TVM_FFI_ICHECK(refAA.as<ObjAA>() != nullptr);
+  TVM_FFI_ICHECK(refAA.as<ObjB>() == nullptr);
+
+  ffi::ObjectRef refB(ffi::make_object<ObjB>());
+  TVM_FFI_ICHECK_EQ(refB->type_index(), ObjB::RuntimeTypeIndex());
+  TVM_FFI_ICHECK(refB.as<ffi::Object>() != nullptr);
+  TVM_FFI_ICHECK(refB.as<ObjBase>() != nullptr);
+  TVM_FFI_ICHECK(refB.as<ObjA>() == nullptr);
+  TVM_FFI_ICHECK(refB.as<ObjAA>() == nullptr);
+  TVM_FFI_ICHECK(refB.as<ObjB>() != nullptr);
+}

tests/cpp/pattern_match_test.cc

@@ -0,0 +1,170 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "../src/arith/pattern_match.h"
+
+#include <gtest/gtest.h>
+#include <tvm/tirx/analysis.h>
+
+TEST(Pattern, Basic) {
+  using namespace tvm;
+  using namespace tvm::tirx;
+  using namespace tvm::arith;
+  tvm::tirx::PrimVar x("x"), y("y"), z("z");
+  PrimExpr scalable_lanes = Mul(Call(PrimType::Int(32), builtin::vscale(), {}), 4);
+  arith::PVar<PrimExpr> px, py, pz;
+  arith::PVar<DLDataType> pt;
+  arith::PVar<PrimExpr> planes;
+  arith::PCallExpr<PVscaleOp> vscale;
+
+  // arithmetics
+  auto r = 1 + (y + 1);
+  TVM_FFI_ICHECK(!(px + (px + px)).Match(r));
+  TVM_FFI_ICHECK(!(px + (py + py)).Match(r));
+  TVM_FFI_ICHECK((px + (py + pz)).Match(r));
+  auto pattern = px + (py + pz);
+  TVM_FFI_ICHECK(pattern.Match(r));
+  {
+    TVM_FFI_ICHECK((px + (py + px)).Match(r));
+    auto rr = (px + py).Eval();
+
+    TVM_FFI_ICHECK(tirx::ExprDeepEqual()(rr, 1 + y));
+    TVM_FFI_ICHECK(tirx::ExprDeepEqual()(px.Eval() + py.Eval(), 1 + y));
+  }
+  {
+    TVM_FFI_ICHECK((px + max(py, px)).Match((x + 1) + max(y, (x + 1))));
+    TVM_FFI_ICHECK(tirx::ExprDeepEqual()(px.Eval(), x + 1));
+  }
+  TVM_FFI_ICHECK(!(px + min(py, px)).Match((x + 1) + max(y, (x + 1))));
+
+  TVM_FFI_ICHECK((px + min(py, px)).Match(z + min(y, z)));
+  TVM_FFI_ICHECK((px + truncdiv(py, px * py)).Match(x + truncdiv(2, x * 2)));
+  TVM_FFI_ICHECK((px - truncmod(py, px * pz)).Match(x - truncmod(2, x * 2)));
+  TVM_FFI_ICHECK((px - floormod(py, px * PConst<PrimExpr>(2))).Match(x - floormod(2, x * 2)));
+
+  // logicals
+  TVM_FFI_ICHECK((px == pz).Match(x == 1));
+  TVM_FFI_ICHECK((px != pz).Match(x != 1));
+  TVM_FFI_ICHECK((px > py).Match(x > y));
+  TVM_FFI_ICHECK((px < py).Match(x < y));
+  TVM_FFI_ICHECK((px <= py).Match(x <= y));
+  TVM_FFI_ICHECK((px >= py).Match(x >= y));
+  TVM_FFI_ICHECK((px >= py && px < pz).Match(x >= y && x < z));
+  TVM_FFI_ICHECK((!(px > py || px != py)).Match(!(x > y || x != y)));
+  {
+    TVM_FFI_ICHECK(select(px >= pz, py, py + pz).Match(tirx::Select((x + 1) >= 1, y, y + 1)));
+    TVM_FFI_ICHECK(tirx::ExprDeepEqual()(px.Eval(), x + 1));
+  }
+  // bit intrinsics
+  {
+    TVM_FFI_ICHECK((px >> pz).Match(x >> 1));
+    TVM_FFI_ICHECK(is_const_int(pz.Eval(), 1));
+  }
+  TVM_FFI_ICHECK(!(px >> pz).Match(x << 1));
+  TVM_FFI_ICHECK((px << pz).Match(x << 1));
+  TVM_FFI_ICHECK((px & pz).Match(x & 1));
+  TVM_FFI_ICHECK((px | pz).Match(x | 1));
+  TVM_FFI_ICHECK((px ^ pz).Match(x ^ 1));
+  TVM_FFI_ICHECK((px - (~(py | (px * pz)))).Match(x - (~(2 | (x * 2)))));
+  // select
+  {
+    TVM_FFI_ICHECK(select(px > pz, py, py + pz).Match(tirx::Select(x > 1, y, y + 1)));
+    TVM_FFI_ICHECK(is_const_int(pz.Eval(), 1));
+  }
+  TVM_FFI_ICHECK(!select(px > pz, py, py + pz).Match(tirx::Select(x > 2, y, y + 1)));
+  TVM_FFI_ICHECK(!select(px > pz, py, py).Match(tirx::Select(x > 2, y, y + 1)));
+  {
+    TVM_FFI_ICHECK(select(px, py, pz).Match(tirx::Select(x > 2, y, y + 1)));
+    TVM_FFI_ICHECK(tirx::ExprDeepEqual()(pz.Eval(), y + 1));
+  }
+  // if_then_else
+  {
+    TVM_FFI_ICHECK(if_then_else(px > pz, py, py + pz).Match(if_then_else(x > 1, y, y + 1)));
+    TVM_FFI_ICHECK(is_const_int(pz.Eval(), 1));
+  }
+  // cast pattern
+  {
+    TVM_FFI_ICHECK(!cast(PConst<DLDataType>(DLDataType{kDLInt, 32, 1}), px)
+                        .Match(tirx::Cast(PrimType::Float(64), x)));
+    TVM_FFI_ICHECK(cast(pt, px).Match(tirx::Cast(PrimType::Float(64), x)));
+    TVM_FFI_ICHECK((pt.Eval() == DLDataType{kDLFloat, 64, 1}));
+    auto zz = cast(pt, px).Eval();
+    TVM_FFI_ICHECK(
+        (cast(pt, px) - cast(pt, py))
+            .Match(tirx::Cast(PrimType::Float(64), x) - tirx::Cast(PrimType::Int(64), x)));
+    auto expr = tirx::Cast(PrimType::Int(32), tirx::Cast(PrimType::Float(64), x));
+    TVM_FFI_ICHECK(!(cast(pt, cast(pt, px))).Match(expr));
+  }
+  // ramp pattern
+  {
+    TVM_FFI_ICHECK(ramp(px, PConst<PrimExpr>(1), planes).Match(tirx::Ramp(x, 1, 10)));
+    TVM_FFI_ICHECK(planes.Eval().as<IntImmNode>()->value == 10);
+    TVM_FFI_ICHECK(ramp(px, PConst<PrimExpr>(1), planes).Match(tirx::Ramp(x, 1, scalable_lanes)));
+    TVM_FFI_ICHECK((vscale * PConst<PrimExpr>(4)).Match(planes.Eval()));
+    TVM_FFI_ICHECK(!ramp(px, PConst<PrimExpr>(1), planes).Match(tirx::Ramp(x, 2, 10)));
+  }
+  // broadcast pattern
+  {
+    TVM_FFI_ICHECK(broadcast(px, planes).Match(tirx::Broadcast(x, 10)));
+    TVM_FFI_ICHECK(planes.Eval().as<IntImmNode>()->value == 10);
+    TVM_FFI_ICHECK(broadcast(px * py, planes).Match(tirx::Broadcast(x * 10, 10)));
+    TVM_FFI_ICHECK(broadcast(px, planes).Match(tirx::Broadcast(x, scalable_lanes)));
+    TVM_FFI_ICHECK((vscale * PConst<PrimExpr>(4)).Match(planes.Eval()));
+  }
+}
+
+TEST(Pattern, IntImm) {
+  using namespace tvm;
+  tirx::PrimVar tx("tx"), ty("ty");
+  arith::PVar<IntImm> c;
+  arith::PVar<tirx::Var> v;
+  {
+    // We can match integer and Var, both of which are
+    // special case container of Expr
+    TVM_FFI_ICHECK((v * c).Match(tx * 3));
+    TVM_FFI_ICHECK_EQ(c.Eval()->value, 3);
+    TVM_FFI_ICHECK((v * 3).Match(tx * 3));
+  }
+  // cannot match c to ty
+  TVM_FFI_ICHECK(!(v * c).Match(tx * ty));
+  // cannot match tx + 1 to v
+  TVM_FFI_ICHECK(!(v * c).Match((tx + 1) * 3));
+}
+
+TEST(Pattern, MatchWithType) {
+  using namespace tvm;
+  // match expr with specified dtype
+  arith::PVarWithDataType<PrimExpr, arith::PConst<DLDataType>> pat(DLDataType{kDLFloat, 32, 1});
+  tirx::PrimVar x("x", PrimType::Float(32));
+  tirx::PrimVar y("y", PrimType::Float(32));
+  tirx::PrimVar x_int("x", PrimType::Int(32));
+  tirx::PrimVar y_int("y", PrimType::Int(32));
+  TVM_FFI_ICHECK(pat.Match(x + y * 2.0f));
+  TVM_FFI_ICHECK(!pat.Match(x_int + y_int * 2));
+
+  // match vectorized expr with specified element dtype
+  arith::PVecDataType vec_ty(DLDataType{kDLFloat, 32, 1});
+  arith::PVarWithDataType<PrimExpr, arith::PVecDataType> vpat(vec_ty);
+  tirx::PrimVar vx("x", PrimType::Float(32, 8));
+  tirx::PrimVar vy("y", PrimType::Float(32, 8));
+  tirx::PrimVar vx_int("x", PrimType::Int(32, 8));
+  tirx::PrimVar vy_int("y", PrimType::Int(32, 8));
+  TVM_FFI_ICHECK(vpat.Match(vx + vy * tirx::Broadcast(2.0f, 8)));
+  TVM_FFI_ICHECK(!vpat.Match(vx_int + vy_int * tirx::Broadcast(2, 8)));
+}

tests/cpp/random_engine_test.cc

@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/random_engine.h>
+
+TEST(RandomEngine, Randomness) {
+  int64_t rand_state = 0;
+
+  tvm::s_tir::LinearCongruentialEngine rng(&rand_state);
+  rng.Seed(0x114514);
+
+  bool covered[100];
+  memset(covered, 0, sizeof(covered));
+  for (int i = 0; i < 100000; i++) {
+    covered[rng() % 100] = true;
+  }
+  for (int i = 0; i < 100; i++) {
+    TVM_FFI_ICHECK(covered[i]);
+  }
+}
+
+TEST(RandomEngine, Reproducibility) {
+  int64_t rand_state_a = 0, rand_state_b = 0;
+  tvm::s_tir::LinearCongruentialEngine rng_a(&rand_state_a), rng_b(&rand_state_b);
+
+  rng_a.Seed(0x23456789);
+  rng_b.Seed(0x23456789);
+
+  for (int i = 0; i < 100000; i++) {
+    TVM_FFI_ICHECK_EQ(rng_a(), rng_b());
+  }
+}
+
+TEST(RandomEngine, Serialization) {
+  int64_t rand_state_a = 0, rand_state_b = 0;
+  tvm::s_tir::LinearCongruentialEngine rng_a(&rand_state_a), rng_b(&rand_state_b);
+
+  rng_a.Seed(0x56728);
+
+  rand_state_b = rand_state_a;
+  for (int i = 0; i < 100000; i++) TVM_FFI_ICHECK_EQ(rng_a(), rng_b());
+
+  for (int i = 0; i < 123456; i++) rng_a();
+
+  rand_state_b = rand_state_a;
+  for (int i = 0; i < 100000; i++) TVM_FFI_ICHECK_EQ(rng_a(), rng_b());
+}

tests/cpp/runtime/logging_test.cc

@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <tvm/runtime/logging.h>
+
+namespace tvm {
+namespace runtime {
+namespace detail {
+namespace {
+
+TEST(TvmLogDebugSettings, Disabled) {
+  TvmLogDebugSettings settings = TvmLogDebugSettings::ParseSpec(nullptr);
+  EXPECT_FALSE(settings.dlog_enabled());
+
+  settings = TvmLogDebugSettings::ParseSpec("");
+  EXPECT_FALSE(settings.dlog_enabled());
+
+  settings = TvmLogDebugSettings::ParseSpec("0");
+  EXPECT_FALSE(settings.dlog_enabled());
+}
+
+TEST(TvmLogDebugSettings, DlogOnly) {
+  TvmLogDebugSettings settings = TvmLogDebugSettings::ParseSpec("1");
+  EXPECT_TRUE(settings.dlog_enabled());
+  EXPECT_FALSE(settings.VerboseEnabled("my/filesytem/src/foo/bar.cc", 0));
+}
+
+TEST(TvmLogDebugSettings, VLogEnabledDefault) {
+  TvmLogDebugSettings settings = TvmLogDebugSettings::ParseSpec("DEFAULT=3");
+  EXPECT_TRUE(settings.dlog_enabled());
+  EXPECT_TRUE(settings.VerboseEnabled("my/filesytem/src/foo/bar.cc", 3));
+  EXPECT_FALSE(settings.VerboseEnabled("my/filesytem/src/foo/bar.cc", 4));
+}
+
+TEST(TvmLogDebugSettings, VLogEnabledComplex) {
+  TvmLogDebugSettings settings =
+      TvmLogDebugSettings::ParseSpec("foo/bar.cc=3,baz.cc=-1,DEFAULT=2,another/file.cc=4");
+  EXPECT_TRUE(settings.dlog_enabled());
+  EXPECT_TRUE(settings.VerboseEnabled("my/filesystem/src/foo/bar.cc", 3));
+  EXPECT_FALSE(settings.VerboseEnabled("my/filesystem/src/foo/bar.cc", 4));
+  EXPECT_TRUE(settings.VerboseEnabled("my/filesystem/src/foo/other.cc", 2));
+  EXPECT_FALSE(settings.VerboseEnabled("my/filesystem/src/foo/other.cc", 3));
+  EXPECT_FALSE(settings.VerboseEnabled("my/filesystem/src/baz.cc", 0));
+}
+
+#define MATCH_THROW(stmt, err_type, matcher)            \
+  try {                                                 \
+    stmt;                                               \
+  } catch (const err_type& e) {                         \
+    EXPECT_THAT(e.what(), matcher);                     \
+  } catch (...) {                                       \
+    EXPECT_FALSE("stmt threw an unexpected exception"); \
+  }
+
+TEST(TvmLogDebugSettings, IllFormed) {
+  MATCH_THROW(
+      TvmLogDebugSettings::ParseSpec("foo/bar.cc=bogus;"), tvm::ffi::Error,
+      ::testing::HasSubstr("TVM_LOG_DEBUG ill-formed at position 11: invalid level: \"bogus;\""));
+
+  MATCH_THROW(TvmLogDebugSettings::ParseSpec("DEFAULT=2;bar/baz.cc=2"), tvm::ffi::Error,
+              ::testing::HasSubstr(
+                  "TVM_LOG_DEBUG ill-formed at position 8: invalid level: \"2;bar/baz.cc=2\""));
+
+  MATCH_THROW(TvmLogDebugSettings::ParseSpec("DEFAULT=2,bar/baz.cc+2"), tvm::ffi::Error,
+              ::testing::HasSubstr("TVM_LOG_DEBUG ill-formed at position 22: expecting "
+                                   "\"=<level>\" after \"bar/baz.cc+2\""));
+}
+
+TEST(TvmLogDebugSettings, SpecPrefix) {
+  TvmLogDebugSettings settings = TvmLogDebugSettings::ParseSpec(
+      "../src/foo/bar.cc=3,src/baz.cc=3,foo/bar/src/another/file.cc=4");
+  EXPECT_TRUE(settings.dlog_enabled());
+  EXPECT_TRUE(settings.VerboseEnabled("my/filesystem/src/foo/bar.cc", 3));
+  EXPECT_TRUE(settings.VerboseEnabled("foo/bar.cc", 3));
+  EXPECT_TRUE(settings.VerboseEnabled("my/filesystem/src/baz.cc", 3));
+  EXPECT_TRUE(settings.VerboseEnabled("baz.cc", 3));
+  EXPECT_TRUE(settings.VerboseEnabled("my/filesystem/src/another/file.cc", 4));
+  EXPECT_TRUE(settings.VerboseEnabled("another/file.cc", 4));
+}
+
+}  // namespace
+}  // namespace detail
+}  // namespace runtime
+}  // namespace tvm

tests/cpp/support/ring_buffer_test.cc

@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "../../../src/support/ring_buffer.h"
+
+#include <gtest/gtest.h>
+
+namespace tvm {
+namespace support {
+namespace {
+
+TEST(RingBuffer, ReadWrite) {
+  RingBuffer buffer;
+  std::vector<int> data = {1, 2, 3, 4};
+  std::vector<int> output;
+
+  buffer.Write(data.data(), data.size() * 4);
+  ASSERT_EQ(buffer.bytes_available(), data.size() * 4);
+
+  output.resize(4);
+  buffer.Read(output.data(), data.size() * 4);
+
+  for (size_t i = 0; i < output.size(); ++i) {
+    ASSERT_EQ(output[i], data[i]);
+  }
+}
+
+TEST(RingBuffer, ReadWithCallback) {
+  RingBuffer buffer;
+  std::vector<int> data = {1, 2, 3, 4};
+  std::vector<int> output;
+
+  buffer.Write(data.data(), data.size() * 4);
+
+  auto callback0 = [](const char* data, size_t size) -> size_t {
+    const int* iptr = reinterpret_cast<const int*>(data);
+    TVM_FFI_ICHECK_EQ(iptr[0], 1);
+    TVM_FFI_ICHECK_EQ(iptr[1], 2);
+    return size;
+  };
+  buffer.ReadWithCallback(callback0, 2 * sizeof(int));
+  auto callback1 = [](const char* data, size_t size) -> size_t {
+    const int* iptr = reinterpret_cast<const int*>(data);
+    TVM_FFI_ICHECK_EQ(iptr[0], 3);
+    TVM_FFI_ICHECK_EQ(iptr[1], 4);
+    return size;
+  };
+  buffer.ReadWithCallback(callback1, 2 * sizeof(int));
+}
+}  // namespace
+}  // namespace support
+}  // namespace tvm

tests/cpp/support_test.cc

@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/runtime/logging.h>
+
+#include "../../src/support/utils.h"
+
+namespace tvm {
+namespace test {
+
+TEST(HashTests, HashStability) {
+  size_t a = 345292;
+  int b = 795620;
+  EXPECT_EQ(::tvm::support::HashCombine(a, b), 2677237020);
+  uint64_t c = 12345;
+  int d = 987654432;
+  EXPECT_EQ(::tvm::support::HashCombine(c, d), 3642871070);
+  size_t e = 1010101;
+  size_t f = 3030303;
+  EXPECT_EQ(::tvm::support::HashCombine(e, f), 2722928432);
+}
+
+TEST(StartsWithTests, Basic) {
+  EXPECT_TRUE(::tvm::support::StartsWith("abc", "abc"));
+  EXPECT_TRUE(::tvm::support::StartsWith("abcd", "abc"));
+  EXPECT_FALSE(::tvm::support::StartsWith("abc", "abcd"));
+}
+
+}  // namespace test
+}  // namespace tvm

tests/cpp/target/virtual_device_test.cc

@@ -0,0 +1,129 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/target/target.h>
+#include <tvm/target/virtual_device.h>
+
+namespace tvm {
+namespace {
+
+TEST(VirtualDevice, Join_Defined) {
+  {
+    Target target_a = Target("cuda");
+    VirtualDevice lhs = VirtualDevice(kDLCUDA, 3);
+    VirtualDevice rhs = VirtualDevice(kDLCUDA, -1, target_a, "global");
+    ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
+    EXPECT_TRUE(actual.operator bool());
+    VirtualDevice expected = VirtualDevice(kDLCUDA, 3, target_a, "global");
+    EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual.value(), expected));
+  }
+  {
+    Target target_a = Target("cuda");
+    VirtualDevice lhs = VirtualDevice(kDLCUDA, -1, target_a, "global");
+    VirtualDevice rhs = VirtualDevice(kDLCUDA, 3);
+    ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
+    EXPECT_TRUE(actual.operator bool());
+    VirtualDevice expected = VirtualDevice(kDLCUDA, 3, target_a, "global");
+    EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual.value(), expected));
+  }
+  {
+    Target target_a = Target("cuda");
+    VirtualDevice lhs = VirtualDevice(kDLCUDA);
+    VirtualDevice rhs = VirtualDevice(kDLCUDA, 2, target_a);
+    ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
+    EXPECT_TRUE(actual.operator bool());
+    VirtualDevice expected = VirtualDevice(kDLCUDA, 2, target_a);
+    EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual.value(), expected));
+  }
+  {
+    Target target_a = Target("cuda");
+    VirtualDevice lhs = VirtualDevice();
+    VirtualDevice rhs = VirtualDevice(kDLCUDA, 3, target_a, "global");
+    ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
+    EXPECT_TRUE(actual.operator bool());
+    VirtualDevice expected = rhs;
+    EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual.value(), expected));
+  }
+}
+
+TEST(VirtualDevice, Join_Undefined) {
+  {
+    VirtualDevice lhs = VirtualDevice(kDLCUDA);
+    VirtualDevice rhs = VirtualDevice(kDLCPU);
+    ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
+    EXPECT_FALSE(actual);
+  }
+  {
+    VirtualDevice lhs = VirtualDevice(kDLCUDA, 3);
+    VirtualDevice rhs = VirtualDevice(kDLCUDA, 4);
+    ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
+    EXPECT_FALSE(actual);
+  }
+  {
+    VirtualDevice lhs = VirtualDevice(kDLCUDA, 3, Target("cuda"));
+    VirtualDevice rhs = VirtualDevice(kDLCUDA, 3, Target("cuda"));
+    ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
+    EXPECT_FALSE(actual);
+  }
+  {
+    VirtualDevice lhs = VirtualDevice(kDLCUDA, 3, Target("cuda"), "local");
+    VirtualDevice rhs = VirtualDevice(kDLCUDA, 3, Target("cuda"), "global");
+    ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
+    EXPECT_FALSE(actual);
+  }
+}
+
+TEST(VirtualDevice, Default) {
+  Target target_a = Target("cuda");
+  VirtualDevice lhs = VirtualDevice(kDLCUDA, -1, Target(), "global");
+  VirtualDevice rhs = VirtualDevice(kDLCUDA, 3, target_a, "local");
+  VirtualDevice actual = VirtualDevice::Default(lhs, rhs);
+  VirtualDevice expected = VirtualDevice(kDLCUDA, 3, target_a, "global");
+  EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual, expected));
+}
+
+TEST(VirtualDevice, Constructor_Invalid) {
+  EXPECT_ANY_THROW(VirtualDevice(kDLCPU, -1, Target("cuda")));
+}
+
+TEST(VirtualDeviceCache, Memoized) {
+  VirtualDeviceCache cache;
+  Target target_a = Target("cuda");
+  Target target_b = Target("llvm");
+  Target target_c = Target("cuda");
+  VirtualDevice virtual_device_a = cache.Make(kDLCUDA, 3, target_a, "local");
+  VirtualDevice virtual_device_b = cache.Make(kDLCPU, 1, target_b, "global");
+
+  EXPECT_EQ(cache.Make(kDLCUDA, 3, target_a, "local"), virtual_device_a);
+  EXPECT_EQ(cache.Make(kDLCPU, 1, target_b, "global"), virtual_device_b);
+  EXPECT_NE(cache.Make(kDLCUDA, 2, target_a, "local"), virtual_device_a);
+  EXPECT_NE(cache.Make(kDLCPU, 3, target_b, "local"), virtual_device_a);
+  EXPECT_NE(cache.Make(kDLCUDA, 3, target_a, "global"), virtual_device_a);
+  EXPECT_EQ(cache.Make(kDLCUDA, 3, Target("cuda"), "local"), virtual_device_a);
+  EXPECT_NE(
+      cache.Make(kDLCUDA, 3,
+                 Target(ffi::Map<ffi::String, ffi::Any>{{"kind", ffi::String("cuda")},
+                                                        {"max_threads_per_block", int64_t(4096)}}),
+                 "local"),
+      virtual_device_a);
+}
+
+}  // namespace
+}  // namespace tvm

tests/cpp/target_test.cc

@@ -0,0 +1,396 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/ir/expr.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/target/target.h>
+
+#include <cmath>
+#include <string>
+
+using namespace tvm;
+
+TVM_REGISTER_TARGET_KIND("TestTargetKind", kDLCPU)
+    .set_attr<std::string>("Attr1", "Value1")
+    .add_attr_option<bool>("my_bool")
+    .add_attr_option<ffi::Array<ffi::String>>("your_names")
+    .add_attr_option<ffi::Map<ffi::String, int64_t>>("her_maps");
+
+ffi::Map<ffi::String, ffi::Any> TestTargetParser(ffi::Map<ffi::String, ffi::Any> target) {
+  ffi::String mcpu = target.at("mcpu").as_or_throw<ffi::String>();
+  target.Set("mcpu", ffi::String("super_") + mcpu);
+  target.Set("keys", ffi::Array<ffi::String>({"super"}));
+  target.Set("feature.test", true);
+  return target;
+}
+
+ffi::Map<ffi::String, ffi::Any> TestAttrsPreProcessor(ffi::Map<ffi::String, ffi::Any> attrs) {
+  attrs.Set("mattr", ffi::String("woof"));
+  return attrs;
+}
+
+TVM_REGISTER_TARGET_KIND("TestTargetParser", kDLCPU)
+    .add_attr_option<ffi::String>("mattr")
+    .add_attr_option<ffi::String>("mcpu")
+    .set_default_keys({"cpu"})
+    .set_target_canonicalizer(TestTargetParser);
+
+TVM_REGISTER_TARGET_KIND("TestAttrsPreprocessor", kDLCPU)
+    .add_attr_option<ffi::String>("mattr")
+    .set_default_keys({"cpu"})
+    .set_target_canonicalizer(TestAttrsPreProcessor);
+
+TVM_REGISTER_TARGET_KIND("TestClashingPreprocessor", kDLCPU)
+    .add_attr_option<ffi::String>("mattr")
+    .add_attr_option<ffi::String>("mcpu")
+    .set_default_keys({"cpu"})
+    .set_target_canonicalizer(TestTargetParser);
+
+TEST(TargetKind, GetAttrMap) {
+  auto map = tvm::TargetKind::GetAttrMap<std::string>("Attr1");
+  auto target_kind = tvm::TargetKind::Get("TestTargetKind").value();
+  std::string result = map[target_kind];
+  TVM_FFI_ICHECK_EQ(result, "Value1");
+}
+
+TEST(TargetCreation, NestedConfig) {
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"my_bool", true},
+      {"your_names", ffi::Array<ffi::String>{"junru", "jian"}},
+      {"kind", ffi::String("TestTargetKind")},
+      {
+          "her_maps",
+          ffi::Map<ffi::String, int64_t>{
+              {"a", 1},
+              {"b", 2},
+          },
+      },
+  };
+  Target target = Target(config);
+  TVM_FFI_ICHECK_EQ(target->kind, TargetKind::Get("TestTargetKind").value());
+  TVM_FFI_ICHECK_EQ(target->tag, "");
+  TVM_FFI_ICHECK(target->keys.empty());
+  bool my_bool = target->GetAttr<bool>("my_bool").value();
+  TVM_FFI_ICHECK_EQ(my_bool, true);
+  ffi::Array<ffi::String> your_names =
+      target->GetAttr<ffi::Array<ffi::String>>("your_names").value();
+  TVM_FFI_ICHECK_EQ(your_names.size(), 2U);
+  TVM_FFI_ICHECK_EQ(your_names[0], "junru");
+  TVM_FFI_ICHECK_EQ(your_names[1], "jian");
+  ffi::Map<ffi::String, int64_t> her_maps =
+      target->GetAttr<ffi::Map<ffi::String, int64_t>>("her_maps").value();
+  TVM_FFI_ICHECK_EQ(her_maps.size(), 2U);
+  TVM_FFI_ICHECK_EQ(her_maps["a"], 1);
+  TVM_FFI_ICHECK_EQ(her_maps["b"], 2);
+}
+
+TEST(TargetCreationFail, UnrecognizedConfigOption) {
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"my_bool", true},
+      {"your_names", ffi::Array<ffi::String>{"junru", "jian"}},
+      {"kind", ffi::String("TestTargetKind")},
+      {"bad", ffi::ObjectRef(nullptr)},
+      {
+          "her_maps",
+          ffi::Map<ffi::String, int64_t>{
+              {"a", 1},
+              {"b", 2},
+          },
+      },
+  };
+  bool failed = false;
+  try {
+    Target tgt(config);
+  } catch (...) {
+    failed = true;
+  }
+  ASSERT_EQ(failed, true);
+}
+
+TEST(TargetCreationFail, TypeMismatch) {
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"my_bool", ffi::String("true")},
+      {"your_names", ffi::Array<ffi::String>{"junru", "jian"}},
+      {"kind", ffi::String("TestTargetKind")},
+      {
+          "her_maps",
+          ffi::Map<ffi::String, int64_t>{
+              {"a", 1},
+              {"b", 2},
+          },
+      },
+  };
+  bool failed = false;
+  try {
+    Target tgt(config);
+  } catch (...) {
+    failed = true;
+  }
+  ASSERT_EQ(failed, true);
+}
+
+TEST(TargetCreationFail, TargetKindNotFound) {
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"my_bool", "true"},
+      {"your_names", ffi::Array<ffi::String>{"junru", "jian"}},
+      {
+          "her_maps",
+          ffi::Map<ffi::String, int64_t>{
+              {"a", 1},
+              {"b", 2},
+          },
+      },
+  };
+  bool failed = false;
+  try {
+    Target tgt(config);
+  } catch (...) {
+    failed = true;
+  }
+  ASSERT_EQ(failed, true);
+}
+
+TEST(TargetCreation, TargetParser) {
+  Target test_target(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestTargetParser")},
+      {"mcpu", ffi::String("woof")},
+  });
+  ASSERT_EQ(test_target->GetAttr<ffi::String>("mcpu").value(), "super_woof");
+  ASSERT_EQ(test_target->keys.size(), 1);
+  ASSERT_EQ(test_target->keys[0], "super");
+}
+
+TEST(TargetCreation, TargetFeatures) {
+  Target test_target_with_parser(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestTargetParser")},
+      {"mcpu", ffi::String("woof")},
+  });
+  // Features are stored as "feature.xxx" keys in attrs
+  ASSERT_EQ(test_target_with_parser->GetAttr<bool>("feature.test").value(), true);
+
+  Target test_target_no_parser("TestTargetKind");
+  ASSERT_EQ(test_target_no_parser->GetAttr<bool>("feature.test"), std::nullopt);
+  ASSERT_EQ(test_target_no_parser->GetAttr<bool>("feature.test", true).value(), true);
+}
+
+TEST(TargetCreation, TargetFeaturesSetByCanonicalizer) {
+  // feature.* keys are set by the canonicalizer, not by user input.
+  Target test_target(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestTargetParser")},
+      {"mcpu", ffi::String("woof")},
+  });
+  // TestTargetParser sets "feature.test" = true
+  ASSERT_EQ(test_target->GetAttr<bool>("feature.test").value(), true);
+  // Non-existent features return nullopt
+  ASSERT_EQ(test_target->GetAttr<bool>("feature.nonexistent"), std::nullopt);
+}
+
+TEST(TargetCreation, TargetAttrsPreProcessor) {
+  Target test_target(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestAttrsPreprocessor")},
+      {"mattr", ffi::String("cake")},
+  });
+  ASSERT_EQ(test_target->GetAttr<ffi::String>("mattr").value(), "woof");
+}
+
+TEST(TargetCreation, TargetParserProcessing) {
+  Target test_target(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestClashingPreprocessor")},
+      {"mcpu", ffi::String("woof")},
+      {"mattr", ffi::String("cake")},
+  });
+  ASSERT_EQ(test_target->GetAttr<ffi::String>("mcpu").value(), "super_woof");
+  ASSERT_EQ(test_target->GetAttr<ffi::String>("mattr").value(), "cake");
+}
+
+TEST(TargetCreation, RoundTripCanonicalizerFeatures) {
+  // Construct a target whose canonicalizer sets feature.test and transforms mcpu
+  Target original(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestTargetParser")},
+      {"mcpu", ffi::String("woof")},
+  });
+  ASSERT_EQ(original->GetAttr<ffi::String>("mcpu").value(), "super_woof");
+  ASSERT_EQ(original->GetAttr<bool>("feature.test").value(), true);
+
+  // Export to config and reconstruct
+  ffi::Map<ffi::String, ffi::Any> exported = original->ToConfig();
+  Target reconstructed(exported);
+
+  // Canonicalized attrs must survive the round-trip
+  // Note: mcpu gets canonicalized again (super_super_woof) because the canonicalizer runs
+  ASSERT_TRUE(reconstructed->GetAttr<ffi::String>("mcpu").has_value());
+  ASSERT_EQ(reconstructed->GetAttr<bool>("feature.test").value(), true);
+  ASSERT_EQ(reconstructed->keys.size(), 1);
+  ASSERT_EQ(reconstructed->keys[0], "super");
+}
+
+TEST(TargetCreation, RoundTripCanonicalizerFeaturesNestedHost) {
+  // Construct a host target whose canonicalizer sets feature.test
+  Target host(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestTargetParser")},
+      {"mcpu", ffi::String("woof")},
+  });
+  ASSERT_EQ(host->GetAttr<bool>("feature.test").value(), true);
+
+  // Attach it as host to another target
+  Target outer(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestTargetKind")},
+      {"my_bool", true},
+  });
+  Target combined(outer, host);
+
+  // Export the outer target (includes nested host) and reconstruct
+  ffi::Map<ffi::String, ffi::Any> exported = combined->ToConfig();
+  Target reconstructed(exported);
+
+  // The nested host must reconstruct successfully with feature.* preserved
+  ffi::Optional<Target> reconstructed_host = reconstructed->GetHost();
+  ASSERT_TRUE(reconstructed_host.has_value());
+  ASSERT_EQ(reconstructed_host.value()->GetAttr<bool>("feature.test").value(), true);
+  ASSERT_TRUE(reconstructed_host.value()->GetAttr<ffi::String>("mcpu").has_value());
+}
+
+TEST(TargetCreationFail, UnknownNonFeatureKeyStillFails) {
+  // Verify that unknown non-feature.* keys still fail schema validation
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"kind", ffi::String("TestTargetParser")},
+      {"mcpu", ffi::String("woof")},
+      {"unknown_key", ffi::String("bad")},
+  };
+  ASSERT_THROW({ Target{config}; }, tvm::ffi::Error);
+}
+
+TVM_REGISTER_TARGET_KIND("TestStringKind", kDLCPU)
+    .add_attr_option<ffi::String>("single")
+    .add_attr_option<ffi::Array<ffi::String>>("array")
+    .add_attr_option<ffi::Array<ffi::Array<ffi::String>>>("nested-array")
+    .add_attr_option<ffi::Array<ffi::Array<ffi::Array<ffi::String>>>>("nested2-array");
+
+TEST(TargetCreation, ProcessStrings) {
+  // Test single string attribute
+  Target test_target1(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestStringKind")},
+      {"single", ffi::String("'string with single quote")},
+  });
+  ASSERT_TRUE(test_target1->GetAttr<ffi::String>("single"));
+  ffi::String string1 = test_target1->GetAttr<ffi::String>("single").value();
+  ASSERT_EQ(string1, "'string with single quote");
+
+  // Test array of strings
+  Target test_target3(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestStringKind")},
+      {"array", ffi::Array<ffi::String>{"-danny", "-sammy=1", "-kirby=string with space"}},
+  });
+  ASSERT_TRUE(test_target3->GetAttr<ffi::Array<ffi::String>>("array"));
+  ffi::Array<ffi::String> array3 = test_target3->GetAttr<ffi::Array<ffi::String>>("array").value();
+  ASSERT_EQ(array3[0], "-danny");
+  ASSERT_EQ(array3[1], "-sammy=1");
+  ASSERT_EQ(array3[2], "-kirby=string with space");
+
+  // Test nested array of strings
+  Target test_target6(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestStringKind")},
+      {"nested-array",
+       ffi::Array<ffi::Array<ffi::String>>{
+           ffi::Array<ffi::String>{"foo0", "foo1", "foo2"},
+           ffi::Array<ffi::String>{"bar0", "bar1", "bar2"},
+           ffi::Array<ffi::String>{"baz0", "baz1"},
+       }},
+  });
+  ASSERT_TRUE(test_target6->GetAttr<ffi::Array<ffi::Array<ffi::String>>>("nested-array"));
+  ffi::Array<ffi::Array<ffi::String>> array6 =
+      test_target6->GetAttr<ffi::Array<ffi::Array<ffi::String>>>("nested-array").value();
+  ASSERT_EQ(array6[0][0], "foo0");
+  ASSERT_EQ(array6[0][1], "foo1");
+  ASSERT_EQ(array6[0][2], "foo2");
+  ASSERT_EQ(array6[1][0], "bar0");
+  ASSERT_EQ(array6[1][1], "bar1");
+  ASSERT_EQ(array6[1][2], "bar2");
+  ASSERT_EQ(array6[2][0], "baz0");
+  ASSERT_EQ(array6[2][1], "baz1");
+
+  // Test doubly-nested array of strings
+  Target test_target7(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestStringKind")},
+      {"nested2-array",
+       ffi::Array<ffi::Array<ffi::Array<ffi::String>>>{
+           ffi::Array<ffi::Array<ffi::String>>{
+               ffi::Array<ffi::String>{"foo0", "foo1"},
+               ffi::Array<ffi::String>{"bar0", "bar1"},
+               ffi::Array<ffi::String>{"baz0", "baz1"},
+           },
+           ffi::Array<ffi::Array<ffi::String>>{
+               ffi::Array<ffi::String>{"zing0", "zing1"},
+               ffi::Array<ffi::String>{"fred"},
+           },
+       }},
+  });
+  ASSERT_TRUE(
+      test_target7->GetAttr<ffi::Array<ffi::Array<ffi::Array<ffi::String>>>>("nested2-array"));
+  ffi::Array<ffi::Array<ffi::Array<ffi::String>>> array7 =
+      test_target7->GetAttr<ffi::Array<ffi::Array<ffi::Array<ffi::String>>>>("nested2-array")
+          .value();
+  ASSERT_EQ(array7.size(), 2);
+  ASSERT_EQ(array7[0].size(), 3);
+  ASSERT_EQ(array7[0][0].size(), 2);
+  ASSERT_EQ(array7[0][1].size(), 2);
+  ASSERT_EQ(array7[0][2].size(), 2);
+  ASSERT_EQ(array7[1].size(), 2);
+  ASSERT_EQ(array7[1][0].size(), 2);
+  ASSERT_EQ(array7[1][1].size(), 1);
+
+  ASSERT_EQ(array7[0][0][0], "foo0");
+  ASSERT_EQ(array7[0][0][1], "foo1");
+  ASSERT_EQ(array7[0][1][0], "bar0");
+  ASSERT_EQ(array7[0][1][1], "bar1");
+  ASSERT_EQ(array7[0][2][0], "baz0");
+  ASSERT_EQ(array7[0][2][1], "baz1");
+  ASSERT_EQ(array7[1][0][0], "zing0");
+  ASSERT_EQ(array7[1][0][1], "zing1");
+  ASSERT_EQ(array7[1][1][0], "fred");
+}
+
+TEST(TargetCreation, DeduplicateKeys) {
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"kind", ffi::String("llvm")},
+      {"keys", ffi::Array<ffi::String>{"cpu", "arm_cpu"}},
+      {"device", ffi::String("arm_cpu")},
+  };
+  Target target = Target(config);
+  TVM_FFI_ICHECK_EQ(target->kind, TargetKind::Get("llvm").value());
+  TVM_FFI_ICHECK_EQ(target->tag, "");
+  TVM_FFI_ICHECK_EQ(target->keys.size(), 2U);
+  TVM_FFI_ICHECK_EQ(target->keys[0], "cpu");
+  TVM_FFI_ICHECK_EQ(target->keys[1], "arm_cpu");
+  TVM_FFI_ICHECK_EQ(target->attrs.count("keys"), 0U);
+  TVM_FFI_ICHECK_EQ(target->GetAttr<ffi::String>("device"), "arm_cpu");
+}
+
+TEST(TargetKindRegistry, ListTargetKinds) {
+  ffi::Array<ffi::String> names = TargetKindRegEntry::ListTargetKinds();
+  TVM_FFI_ICHECK_EQ(names.empty(), false);
+  TVM_FFI_ICHECK_EQ(std::count(std::begin(names), std::end(names), "llvm"), 1);
+}
+
+TEST(TargetKindRegistry, ListTargetOptions) {
+  TargetKind llvm = TargetKind::Get("llvm").value();
+  ffi::Map<ffi::String, ffi::String> attrs = TargetKindRegEntry::ListTargetKindOptions(llvm);
+  TVM_FFI_ICHECK_EQ(attrs.empty(), false);
+}

tests/cpp/te_compute_test.cc

@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/te/operation.h>
+
+TEST(Tensor, Basic) {
+  using namespace tvm;
+  using namespace tvm::te;
+
+  PrimVar m("m"), n("n"), l("l");
+
+  Tensor A = placeholder({m, l}, PrimType::Float(32), "A");
+  Tensor B = placeholder({n, l}, PrimType::Float(32), "B");
+
+  auto C = compute({m, n}, [&](PrimVar i, PrimVar j) { return A[i][j]; }, "C");
+
+  Tensor::Slice x = A[n];
+}
+
+TEST(Tensor, Reduce) {
+  using namespace tvm;
+  using namespace tvm::te;
+
+  PrimVar m("m"), n("n"), l("l");
+  te::Tensor A = te::placeholder({m, l}, PrimType::Float(32), "A");
+  te::Tensor B = te::placeholder({n, l}, PrimType::Float(32), "B");
+  IterVar rv = reduce_axis(Range{0, l}, "k");
+
+  auto C = te::compute(
+      {m, n}, [&](PrimVar i, PrimVar j) { return sum(max(1 + A[i][rv] + 1, B[j][rv]), {rv}); },
+      "C");
+}
+
+TEST(Tensor, Indexing) {
+  using namespace tvm;
+  using namespace tvm::te;
+
+  PrimVar x("x"), y("y");
+  te::Tensor A = te::placeholder({x, y}, PrimType::Float(32), "A");
+}

tests/cpp/threading_backend_test.cc

@@ -0,0 +1,196 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "../../src/runtime/threading_backend.h"
+
+#include <gtest/gtest.h>
+#include <tvm/runtime/c_backend_api.h>
+#include <tvm/runtime/logging.h>
+
+#include <atomic>
+#include <memory>
+#include <mutex>
+#include <sstream>
+#include <thread>
+#include <unordered_map>
+#include <unordered_set>
+
+constexpr size_t N = 128;
+void AtomicCompute(int task_id, size_t n, std::atomic<size_t>* acc, TVMParallelGroupEnv* penv) {
+  const size_t N_per_task = (n + penv->num_task - 1) / penv->num_task;
+  for (size_t i = task_id * N_per_task; i < n && i < (task_id + 1) * N_per_task; ++i) {
+    acc->fetch_add(i, std::memory_order_relaxed);
+  }
+  return;
+}
+
+class AffinityCheck {
+ public:
+  AffinityCheck(uint32_t parent_id, int max_concurrency, std::atomic<size_t>* acc)
+      : id_(parent_id), max_concurrency_(max_concurrency), acc_(acc) {}
+
+  void Compute(int task_id, size_t n, TVMParallelGroupEnv* penv) {
+    AtomicCompute(task_id, n, acc_, penv);
+  }
+
+  int GetComputeResult() { return acc_->load(std::memory_order_relaxed); }
+
+  void GetAffinity(int task_id) {
+#if defined(__linux__)
+    cpu_set_t cpuset;
+    CPU_ZERO(&cpuset);
+    pthread_getaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
+    std::lock_guard<std::mutex> lock(mutex_);
+    thread_affinity_[task_id] = cpuset;
+    // Printing the current thread CPU affinity.
+    std::ostringstream str;
+    for (int i = 0; i < max_concurrency_; i++) {
+      if (CPU_ISSET(i, &cpuset)) {
+        str << i << ",";
+      }
+    }
+#endif
+  }
+
+  bool VerifyAffinity(const std::vector<uint32_t>& cpus) {
+#if defined(__linux__)
+    std::unordered_set<uint32_t> uset;
+    cpu_set_t cpu_mask;
+    CPU_ZERO(&cpu_mask);
+    for (auto x : cpus) {
+      CPU_SET(x, &cpu_mask);
+      uset.insert(x);
+    }
+
+    for (auto x : thread_affinity_) {
+      if (!CPU_EQUAL(&cpu_mask, &x.second)) {
+        bool cpu_find = false;
+        for (auto cpu : uset) {
+          CPU_ISSET(cpu, &x.second);
+          uset.erase(cpu);
+          cpu_find = true;
+          break;
+        }
+        if (!cpu_find) return false;
+      }
+    }
+#endif
+    return true;
+  }
+
+ private:
+  uint32_t id_;
+  int max_concurrency_;
+  std::atomic<size_t>* acc_;
+  std::mutex mutex_;
+#if defined(__linux__)
+  std::unordered_map<int, cpu_set_t> thread_affinity_;
+#endif
+};
+
+static FTVMParallelLambda atomic_add_task_id = [](int task_id, TVMParallelGroupEnv* penv,
+                                                  void* cdata) -> int {
+  auto* data = reinterpret_cast<std::atomic<size_t>*>(cdata);
+  AtomicCompute(task_id, N, data, penv);
+  return 0;
+};
+
+static FTVMParallelLambda affinity_check_task_id = [](int task_id, TVMParallelGroupEnv* penv,
+                                                      void* cdata) -> int {
+  auto* data = reinterpret_cast<AffinityCheck*>(cdata);
+  data->Compute(task_id, N, penv);
+  data->GetAffinity(task_id);
+  return 0;
+};
+
+TEST(ThreadingBackend, TVMBackendParallelLaunch) {
+  std::atomic<size_t> acc(0);
+  TVMBackendParallelLaunch(atomic_add_task_id, &acc, 0);
+  EXPECT_EQ(acc.load(std::memory_order_relaxed), N * (N - 1) / 2);
+}
+
+TEST(ThreadingBackend, TVMBackendParallelLaunchMultipleThreads) {
+  // TODO(tulloch) use parameterised tests when available.
+  size_t num_jobs_per_thread = 3;
+  size_t max_num_threads = 2;
+
+  for (size_t num_threads = 1; num_threads < max_num_threads; ++num_threads) {
+    std::vector<std::unique_ptr<std::thread>> ts;
+    for (size_t i = 0; i < num_threads; ++i) {
+      ts.emplace_back(new std::thread([&]() {
+        for (size_t j = 0; j < num_jobs_per_thread; ++j) {
+          std::atomic<size_t> acc(0);
+          TVMBackendParallelLaunch(atomic_add_task_id, &acc, 0);
+          EXPECT_EQ(acc.load(std::memory_order_relaxed), N * (N - 1) / 2);
+        }
+      }));
+    }
+    for (auto& t : ts) {
+      t->join();
+    }
+  }
+}
+
+TEST(ThreadingBackend, TVMBackendAffinityConfigure) {
+  int max_concurrency = tvm::runtime::threading::MaxConcurrency();
+  std::vector<std::unique_ptr<std::thread>> ts;
+  // Returning as there is only one CPU available.
+  if (max_concurrency <= 1) {
+    return;
+  }
+  // Creating two threads to test the 'CPU list affinity' feature.
+  const int threads_num = 2;
+  // Getting the maximum number of CPUs which are available to each thread.
+  const int cpus_num_per_thread = max_concurrency / threads_num;
+  // Testing two mode of affinity.,
+  std::vector<tvm::runtime::threading::ThreadGroup::AffinityMode> modes = {
+      tvm::runtime::threading::ThreadGroup::kSpecifyOneCorePerThread,
+      tvm::runtime::threading::ThreadGroup::kSpecifyThreadShareAllCore};
+  for (auto mode : modes) {
+    for (int thread_pool_idx = 0; thread_pool_idx < threads_num; thread_pool_idx++) {
+      ts.emplace_back(new std::thread(
+          [&](int thread_pool_index, int sys_max_concurrency,
+              tvm::runtime::threading::ThreadGroup::AffinityMode affinity_mode) {
+            std::atomic<size_t> acc(0);
+            AffinityCheck ac(thread_pool_index, sys_max_concurrency, &acc);
+            std::vector<unsigned int> cpus;
+            for (int k = 0; k < cpus_num_per_thread; k++) {
+              cpus.push_back(thread_pool_index * cpus_num_per_thread + k);
+            }
+            tvm::runtime::threading ::Configure(affinity_mode, 0, cpus);
+            TVMBackendParallelLaunch(affinity_check_task_id, &ac, 0);
+            EXPECT_EQ(ac.GetComputeResult(), N * (N - 1) / 2);
+            EXPECT_EQ(ac.VerifyAffinity(cpus), true);
+          },
+          thread_pool_idx, max_concurrency, mode));
+    }
+  }
+  for (auto& t : ts) {
+    t->join();
+  }
+}
+
+TEST(ThreadingBackend, TVMBackendParallelForWithThreadingBackend) {
+  int n = 100;
+  std::vector<int> vec(/*size=*/n, /*value=*/0);
+  tvm::runtime::parallel_for_with_threading_backend([&vec](int i) { vec[i] = i; }, 0, n);
+  for (int i = 0; i < n; ++i) {
+    EXPECT_EQ(vec[i], i);
+  }
+}

tests/cpp/tir_analysis_side_effect.cc

@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/te/operation.h>
+#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/builtin.h>
+
+TEST(SimplePasses, SideEffect) {
+  using namespace tvm;
+  auto buf = tirx::decl_buffer({16}, PrimType::Float(32));
+  auto i = tirx::PrimVar("i", PrimType::Int(32));
+  TVM_FFI_ICHECK(tirx::SideEffect(tirx::BufferLoad(buf, {i})) == tirx::CallEffectKind::kReadState);
+  TVM_FFI_ICHECK(tirx::SideEffect(exp(tirx::Cast(PrimType::Float(32), i + 1))) ==
+                 tirx::CallEffectKind::kPure);
+  TVM_FFI_ICHECK(tirx::SideEffect(tvm::Call(PrimType::Void(), tirx::builtin::tvm_storage_sync(), {})
+                                      .as_or_throw<PrimExpr>()) ==
+                 tirx::CallEffectKind::kUpdateState);
+}

tests/cpp/tir_scalable_datatype.cc

@@ -0,0 +1,212 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <tvm/ffi/dtype.h>
+#include <tvm/script/printer/printer.h>
+#include <tvm/tirx/builtin.h>
+#include <tvm/tirx/expr.h>
+
+#ifdef TVM_LLVM_VERSION
+#include <llvm/IR/Intrinsics.h>
+#endif
+
+#include "../../src/script/printer/utils.h"
+
+using ::testing::HasSubstr;
+
+// ---------
+// Prim Type
+// ---------
+TEST(ScalablePrimType, TestCreateScalableType) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLInt, 32, 4);
+  ASSERT_EQ(scalable_type.code(), kDLInt);
+  ASSERT_EQ(scalable_type.bits(), 32);
+  ASSERT_EQ(scalable_type.VScaleFactor(), 4);
+  ASSERT_TRUE(scalable_type.IsScalableVector());
+  ASSERT_TRUE(scalable_type.IsScalableVector() || scalable_type.IsFixedLengthVector());
+}
+
+TEST(ScalablePrimType, TestScalableWithBits) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLInt, 1, 8);
+  scalable_type = scalable_type.WithBits(32);
+  ASSERT_EQ(scalable_type.bits(), 32);
+  ASSERT_TRUE(scalable_type.IsScalableVector());
+  ASSERT_TRUE(scalable_type.IsScalableVector() || scalable_type.IsFixedLengthVector());
+}
+
+TEST(ScalablePrimType, TestScalableWithVscaleFactor) {
+  tvm::PrimType type = tvm::PrimType::Int(32);
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(type.code(), type.bits(), 4);
+  ASSERT_EQ(scalable_type.VScaleFactor(), 4);
+  ASSERT_TRUE(scalable_type.IsScalableVector());
+  ASSERT_TRUE(scalable_type.IsScalableVector() || scalable_type.IsFixedLengthVector());
+}
+
+TEST(ScalablePrimType, TestAssignScalablePrimType) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLInt, 32, 2);
+  tvm::PrimType scalable_type_copy = scalable_type;
+  ASSERT_TRUE(scalable_type_copy.IsScalableVector());
+  ASSERT_TRUE(scalable_type_copy.IsScalableVector() || scalable_type_copy.IsFixedLengthVector());
+}
+
+TEST(ScalablePrimType, TestScalablePrimTypeEquality) {
+  ASSERT_TRUE(tvm::PrimType::ScalableVector(kDLInt, 32, 4) ==
+              tvm::PrimType::ScalableVector(kDLInt, 32, 4));
+}
+
+TEST(ScalablePrimType, TestScalablePrimTypeAndNonScalablePrimTypeInequality) {
+  ASSERT_FALSE(tvm::PrimType::ScalableVector(kDLInt, 32, 4) == tvm::PrimType::Int(32, 4));
+}
+
+TEST(ScalablePrimType, TestIsScalar) {
+  ASSERT_FALSE(tvm::PrimType::ScalableVector(kDLInt, 32, 4).IsScalar());
+  ASSERT_TRUE(tvm::PrimType::Int(32).IsScalar());
+  ASSERT_FALSE(tvm::PrimType::Int(32, 4).IsScalar());
+  ASSERT_FALSE(tvm::PrimType::Void().IsScalar());
+}
+
+TEST(ScalablePrimType, TestScalablePrimTypeToString) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLInt, 32, 4);
+  EXPECT_EQ(tvm::ffi::DLDataTypeToString(scalable_type->dtype), "int32xvscalex4");
+}
+
+TEST(ScalablePrimType, TestStringToScalablePrimType) {
+  std::string scalable_type_str = "int32xvscalex4";
+  EXPECT_EQ(tvm::PrimType(tvm::ffi::StringToDLDataType(scalable_type_str)),
+            tvm::PrimType::ScalableVector(kDLInt, 32, 4));
+}
+
+TEST(ScalablePrimType, TestInvalidStringToScalablePrimType) {
+  std::string scalable_type_str = "int32x4xvscale";
+  EXPECT_THROW(
+      {
+        try {
+          tvm::ffi::StringToDLDataType(scalable_type_str);
+        } catch (const tvm::ffi::Error& e) {
+          EXPECT_THAT(e.what(), HasSubstr("unknown dtype `int32x4xvscale`"));
+          throw;
+        }
+      },
+      tvm::ffi::Error);
+}
+
+TEST(ScalablePrimType, TestGetScalableVectorBytes) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLInt, 32, 4);
+  EXPECT_THROW(
+      {
+        try {
+          int bytes = (scalable_type.bits() * scalable_type.lanes() + 7) / 8;
+          static_cast<void>(bytes);
+        } catch (const tvm::ffi::Error& e) {
+          EXPECT_THAT(e.what(),
+                      HasSubstr("Can't fetch the lanes of a scalable vector at a compile time"));
+          throw;
+        }
+      },
+      tvm::ffi::Error);
+}
+
+TEST(ScalablePrimType, TestScalablePrimTypeInvalidLanesError) {
+  EXPECT_THROW(
+      {
+        try {
+          tvm::PrimType::ScalableVector(kDLFloat, 62, 1);
+        } catch (const tvm::ffi::Error& e) {
+          EXPECT_THAT(e.what(), HasSubstr("Invalid value for vscale factor"));
+          throw;
+        }
+      },
+      tvm::ffi::Error);
+}
+
+TEST(ScalablePrimType, TestScalablePrimTypeInvalidVscaleFactorAccess) {
+  tvm::PrimType fixed_length_type = tvm::PrimType::Float(32, 4);
+  ASSERT_TRUE(fixed_length_type.IsFixedLengthVector());
+  ASSERT_TRUE(fixed_length_type.IsScalableVector() || fixed_length_type.IsFixedLengthVector());
+  EXPECT_THROW(
+      {
+        try {
+          fixed_length_type.VScaleFactor();
+        } catch (const tvm::ffi::Error& e) {
+          EXPECT_THAT(e.what(), HasSubstr("A fixed length vector doesn't have a vscale factor"));
+          throw;
+        }
+      },
+      tvm::ffi::Error);
+}
+
+TEST(ScalablePrimType, TestScalablePrimTypeInvalidLanesAccess) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLFloat, 32, 4);
+  EXPECT_THROW(
+      {
+        try {
+          scalable_type.lanes();
+        } catch (const tvm::ffi::Error& e) {
+          EXPECT_THAT(e.what(),
+                      HasSubstr("Can't fetch the lanes of a scalable vector at a compile time"));
+          throw;
+        }
+      },
+      tvm::ffi::Error);
+}
+
+TEST(ScalablePrimType, TestScalableBool) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLBool, 8, 4);
+  ASSERT_EQ(scalable_type.code(), kDLBool);
+  ASSERT_EQ(scalable_type.bits(), 8);
+  ASSERT_EQ(scalable_type.VScaleFactor(), 4);
+  ASSERT_TRUE(scalable_type.IsScalableVector());
+}
+
+TEST(ScalablePrimType, TestScalableUInt) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLUInt, 1, 4);
+  ASSERT_EQ(scalable_type.code(), kDLUInt);
+  ASSERT_EQ(scalable_type.bits(), 1);
+  ASSERT_EQ(scalable_type.VScaleFactor(), 4);
+  ASSERT_TRUE(scalable_type.IsScalableVector());
+}
+
+// -----------
+// Integration
+// -----------
+#ifdef TVM_LLVM_VERSION
+TEST(ScalablePrimType, TestScalableIntrinCall) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLInt, 32, 4);
+  tvm::Call call = tvm::Call(scalable_type, tvm::tirx::builtin::call_llvm_intrin(),
+#if TVM_LLVM_VERSION >= 200
+                             {tvm::IntImm::Int32(::llvm::Intrinsic::stepvector)});
+#else
+                             {tvm::IntImm::Int32(::llvm::Intrinsic::experimental_stepvector)});
+#endif
+  ASSERT_EQ(call->ty.as_or_throw<tvm::PrimType>(), scalable_type);
+  ASSERT_EQ(tvm::Script(call),
+#if TVM_LLVM_VERSION >= 200
+            "T.call_llvm_intrin(\"int32xvscalex4\", \"llvm.stepvector\")");
+#else
+            "T.call_llvm_intrin(\"int32xvscalex4\", \"llvm.experimental.stepvector\")");
+#endif
+}
+#endif
+
+TEST(ScalablePrimType, TestTIRScriptScalableDtype2Str) {
+  tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLInt, 32, 4);
+  ASSERT_EQ(tvm::script::printer::DType2Str(scalable_type->dtype), "int32xvscalex4");
+}

tests/cpp/topi_ewise_test.cc

@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/te/operation.h>
+#include <tvm/topi/elemwise.h>
+
+namespace tvm {
+namespace topi {
+TEST(Tensor, Basic) {
+  using namespace tvm;
+  PrimVar m("m"), l("l");
+  Tensor A = placeholder({m, l}, PrimType::Float(32), "A");
+  auto C = topi::exp(A);
+}
+}  // namespace topi
+}  // namespace tvm

tests/lint/check_asf_header.py

@@ -0,0 +1,489 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E741
+"""Helper tool to check and fix ASF headers in source files tracked by git.
+
+Ported from tvm-ffi. Replaces the previous RAT-based check_asf_header.sh
+(which required Java) with a self-contained Python implementation.
+"""
+
+from __future__ import annotations
+
+import argparse
+import fnmatch
+import subprocess
+import sys
+from pathlib import Path
+
+header_cstyle = """
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+""".strip()
+
+header_mdstyle = """
+<!--- Licensed to the Apache Software Foundation (ASF) under one -->
+<!--- or more contributor license agreements.  See the NOTICE file -->
+<!--- distributed with this work for additional information -->
+<!--- regarding copyright ownership.  The ASF licenses this file -->
+<!--- to you under the Apache License, Version 2.0 (the -->
+<!--- "License"); you may not use this file except in compliance -->
+<!--- with the License.  You may obtain a copy of the License at -->
+
+<!---   http://www.apache.org/licenses/LICENSE-2.0 -->
+
+<!--- Unless required by applicable law or agreed to in writing, -->
+<!--- software distributed under the License is distributed on an -->
+<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
+<!--- KIND, either express or implied.  See the License for the -->
+<!--- specific language governing permissions and limitations -->
+<!--- under the License. -->
+""".strip()
+
+header_pystyle = """
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+""".strip()
+
+header_rststyle = """
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+""".strip()
+
+header_groovystyle = """
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+""".strip()
+
+header_cmdstyle = """
+:: Licensed to the Apache Software Foundation (ASF) under one
+:: or more contributor license agreements.  See the NOTICE file
+:: distributed with this work for additional information
+:: regarding copyright ownership.  The ASF licenses this file
+:: to you under the Apache License, Version 2.0 (the
+:: "License"); you may not use this file except in compliance
+:: with the License.  You may obtain a copy of the License at
+::
+::   http://www.apache.org/licenses/LICENSE-2.0
+::
+:: Unless required by applicable law or agreed to in writing,
+:: software distributed under the License is distributed on an
+:: "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+:: KIND, either express or implied.  See the License for the
+:: specific language governing permissions and limitations
+:: under the License.
+""".strip()
+
+FMT_MAP = {
+    "sh": header_pystyle,
+    "cc": header_cstyle,
+    "c": header_cstyle,
+    "cu": header_cstyle,
+    "cuh": header_cstyle,
+    "mm": header_cstyle,
+    "m": header_cstyle,
+    "go": header_cstyle,
+    "java": header_cstyle,
+    "h": header_cstyle,
+    "pyi": header_pystyle,
+    "py": header_pystyle,
+    "pyx": header_pystyle,
+    "toml": header_pystyle,
+    "yml": header_pystyle,
+    "yaml": header_pystyle,
+    "rs": header_cstyle,
+    "md": header_mdstyle,
+    "cmake": header_pystyle,
+    "mk": header_pystyle,
+    "rst": header_rststyle,
+    "gradle": header_groovystyle,
+    "groovy": header_groovystyle,
+    "tcl": header_pystyle,
+    "xml": header_mdstyle,
+    "storyboard": header_mdstyle,
+    "pbxproj": header_cstyle,
+    "plist": header_mdstyle,
+    "xcworkspacedata": header_mdstyle,
+    "html": header_mdstyle,
+    "bat": header_cmdstyle,
+}
+
+# Files and patterns to skip during header checking.
+# Files and patterns to skip (3rdparty, generated files, etc.).
+SKIP_LIST: list[str] = [
+    "3rdparty/*",
+    "ffi/3rdparty/*",
+    ".github/*",
+    ".txdev/*",
+    ".agents/*",
+    "*.json",
+    "*.txt",
+    "*.svg",
+    "*.lst",
+    "*.lds",
+    "*.in",
+    "*.diff",
+    "*.edl",
+    "*.md5",
+    "*.csv",
+    "*.log",
+    "*.interp",
+    "*.tokens",
+    "*.ipynb",
+    "*.conf",
+    "*.ini",
+    "*.lock",
+    "*.properties",
+    "*.j2",
+]
+
+
+def should_skip_file(filepath: str) -> bool:
+    """Check if file should be skipped based on SKIP_LIST."""
+    for pattern in SKIP_LIST:
+        if fnmatch.fnmatch(filepath, pattern):
+            return True
+    return False
+
+
+def get_git_files() -> list[str] | None:
+    """Get list of files tracked by git (excluding files staged for deletion)."""
+    try:
+        result = subprocess.run(
+            ["git", "ls-files"], check=False, capture_output=True, text=True, cwd=Path.cwd()
+        )
+        if result.returncode != 0:
+            print("Error: Could not get git files. Make sure you're in a git repository.")
+            print("Git command failed:", result.stderr.strip())
+            return None
+        all_files = {line.strip() for line in result.stdout.split("\n") if line.strip()}
+        # Exclude files staged for deletion so the header check does not
+        # report errors for files that are intentionally being removed.
+        deleted_result = subprocess.run(
+            ["git", "ls-files", "--deleted"],
+            check=False,
+            capture_output=True,
+            text=True,
+            cwd=Path.cwd(),
+        )
+        if deleted_result.returncode == 0:
+            deleted = {line.strip() for line in deleted_result.stdout.split("\n") if line.strip()}
+            all_files -= deleted
+        elif deleted_result.stderr:
+            print(f"Warning: 'git ls-files --deleted' failed: {deleted_result.stderr.strip()}")
+        return sorted(all_files)
+    except FileNotFoundError:
+        print("Error: Git not found. This tool requires git to be installed.")
+        return None
+
+
+def copyright_line(line: str) -> bool:
+    # Following two items are intentionally break apart
+    # so that the copyright detector won"t detect the file itself.
+    if line.find("Copyright " + "(c)") != -1:
+        return True
+    # break pattern into two lines to avoid false-negative check
+    spattern1 = "Copyright"
+    if line.find(spattern1) != -1 and line.find("by") != -1:
+        return True
+    return False
+
+
+def check_header(fname: str, header: str) -> bool:
+    """Check header status of file without modifying it."""
+    if not Path(fname).exists():
+        print(f"ERROR: Cannot find {fname}")
+        return False
+
+    lines = Path(fname).open().readlines()
+
+    has_asf_header = False
+    has_copyright = False
+
+    for i, l in enumerate(lines):
+        if l.find("Licensed to the Apache Software Foundation") != -1:
+            has_asf_header = True
+        elif copyright_line(l):
+            has_copyright = True
+
+    if has_asf_header and not has_copyright:
+        return True  # File is good
+
+    if not has_asf_header:
+        print(f"ERROR: Missing ASF header in {fname}")
+        print("Run `python tests/lint/check_asf_header.py --fix` to add the header")
+        return False
+
+    if has_copyright:
+        print(f"ERROR: Has copyright line that should be removed in {fname}")
+        return False
+
+    return True
+
+
+def collect_files() -> list[str] | None:
+    """Collect all files that need header checking from git."""
+    files = []
+
+    # Get files from git (required)
+    git_files = get_git_files()
+
+    if git_files is None:
+        # Git is required, so exit if we can't get files
+        return None
+
+    # Filter git files based on supported types and skip list
+    for git_file in git_files:
+        # Check if file should be skipped
+        if should_skip_file(git_file):
+            continue
+
+        # Check if this file type is supported
+        suffix = git_file.split(".")[-1] if "." in git_file else ""
+        basename = Path(git_file).name
+
+        if (
+            suffix in FMT_MAP
+            or basename == "gradle.properties"
+            or (suffix == "" and basename in ["CMakeLists", "Makefile"])
+        ):
+            files.append(git_file)
+
+    return files
+
+
+def add_header(fname: str, header: str) -> None:
+    """Add header to file."""
+    if not Path(fname).exists():
+        print(f"Cannot find {fname} ...")
+        return
+
+    lines = Path(fname).open().readlines()
+
+    has_asf_header = False
+    has_copyright = False
+
+    for i, l in enumerate(lines):
+        if l.find("Licensed to the Apache Software Foundation") != -1:
+            has_asf_header = True
+        elif copyright_line(l):
+            has_copyright = True
+            lines[i] = ""
+
+    if has_asf_header and not has_copyright:
+        return
+
+    with Path(fname).open("w") as outfile:
+        skipline = False
+        if not lines:
+            skipline = False  # File is empty
+        elif lines[0][:2] == "#!":
+            skipline = True
+        elif lines[0][:2] == "<?":
+            skipline = True
+        elif lines[0].startswith("<html>"):
+            skipline = True
+        elif lines[0].startswith("// !$"):
+            skipline = True
+
+        if skipline:
+            outfile.write(lines[0])
+            if not has_asf_header:
+                outfile.write(header + "\n\n")
+            outfile.write("".join(lines[1:]))
+        else:
+            if not has_asf_header:
+                outfile.write(header + "\n\n")
+            outfile.write("".join(lines))
+    if not has_asf_header:
+        print(f"Add header to {fname}")
+    if has_copyright:
+        print(f"Removed copyright line from {fname}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Check and fix ASF headers in source files tracked by git",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+This tool processes all files tracked by git (using 'git ls-files') and automatically
+skips files matching patterns in SKIP_LIST (build artifacts, third-party files, etc.).
+
+One of --check, --fix, or --dry-run must be specified.
+
+Examples:
+  # Show all files that would be processed (dry run)
+  python check_asf_header.py --dry-run
+
+  # Check all git-tracked files (report errors)
+  python check_asf_header.py --check
+
+  # Fix all git-tracked files (add/update headers)
+  python check_asf_header.py --fix
+        """,
+    )
+
+    parser.add_argument(
+        "--check",
+        action="store_true",
+        help="Check mode: report errors without modifying files",
+    )
+
+    parser.add_argument(
+        "--fix",
+        action="store_true",
+        help="Fix mode: fix files with missing or incorrect headers (default)",
+    )
+
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Dry run: show files that would be processed without doing anything",
+    )
+
+    args = parser.parse_args()
+
+    # Validate arguments
+    if args.check and args.fix:
+        print("Error: Cannot specify both --check and --fix")
+        return 1
+
+    if not args.check and not args.fix and not args.dry_run:
+        print("Error: Must specify one of --check, --fix, or --dry-run")
+        return 1
+
+    # Always use git-based approach
+    files = collect_files()
+
+    if files is None:
+        return 1  # Error already printed by get_git_files()
+
+    if not files:
+        print("No files found to process")
+        return 0
+
+    # Handle dry-run mode
+    if args.dry_run:
+        print(f"Would process {len(files)} files:")
+        for fname in sorted(files):
+            print(f"  {fname}")
+        return 0
+
+    # Determine mode
+    check_only = args.check
+
+    error_count = 0
+    processed_count = 0
+
+    for fname in files:
+        processed_count += 1
+        suffix = fname.split(".")[-1] if "." in fname else ""
+        basename = Path(fname).name
+
+        # Determine header type
+        if suffix in FMT_MAP:
+            header = FMT_MAP[suffix]
+        elif basename == "gradle.properties":
+            header = FMT_MAP["h"]
+        elif suffix == "" and basename in ["CMakeLists", "Makefile"]:
+            header = FMT_MAP["cmake"] if basename == "CMakeLists" else FMT_MAP["mk"]
+        else:
+            if check_only:
+                print(f"ERROR: Cannot handle file type for {fname}")
+                error_count += 1
+            else:
+                print(f"Cannot handle {fname} ...")
+            continue
+
+        if check_only:
+            # Check mode
+            if not check_header(fname, header):
+                error_count += 1
+        else:
+            # Fix mode
+            add_header(fname, header)
+
+    if check_only:
+        if error_count > 0:
+            print(f"\nFound {error_count} errors in {processed_count} files")
+            return 1
+        else:
+            print(f"\nAll {processed_count} files are compliant")
+            return 0
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/lint/check_file_type.py

@@ -0,0 +1,160 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Helper tool to check file types that are allowed to checkin."""
+
+import subprocess
+import sys
+from pathlib import Path
+
+# List of file types we allow
+ALLOW_EXTENSION = {
+    # source code
+    "cc",
+    "c",
+    "h",
+    "S",
+    "cu",
+    "cuh",
+    "m",
+    "mm",
+    "java",
+    "gradle",
+    "groovy",
+    "js",
+    "cjs",
+    "mjs",
+    "ts",
+    "sh",
+    "bat",
+    "py",
+    # configurations
+    "cfg",
+    "mk",
+    "in",
+    "cmake",
+    "xml",
+    "toml",
+    "yml",
+    "yaml",
+    "json",
+    "cfg",
+    "ini",
+    # docs
+    "txt",
+    "md",
+    "rst",
+    "css",
+    "html",
+    "ipynb",
+    # ios
+    "pbxproj",
+    "plist",
+    "xcworkspacedata",
+    "storyboard",
+    "xcscheme",
+    # interface definition
+    "idl",
+    # Jinja2 templates
+    "j2",
+    # images
+    "png",
+    # misc
+    "properties",
+    "template",
+}
+
+# List of file names allowed
+ALLOW_FILE_NAME = {
+    ".gitignore",
+    ".gitattributes",
+    ".gitmodules",
+    ".clang-format",
+    "README",
+    "Makefile",
+    "Doxyfile",
+    "CODEOWNERSHIP",
+    "condarc",
+    "with_the_same_user",
+}
+
+# List of specific files allowed in relpath to <proj_root>
+ALLOW_SPECIFIC_FILE = {"LICENSE", "NOTICE", "KEYS"}
+
+
+def filename_allowed(name: str) -> bool:
+    """Check if name is allowed by the current policy.
+
+    Parameters
+    ----------
+    name : str
+        Input name
+
+    Returns
+    -------
+    allowed : bool
+        Whether the filename is allowed.
+    """
+    arr = name.rsplit(".", 1)
+    if arr[-1] in ALLOW_EXTENSION:
+        return True
+
+    basename = Path(name).name
+    if basename in ALLOW_FILE_NAME:
+        return True
+
+    # Dockerfile and variants like Dockerfile.ci_gpu
+    if basename == "Dockerfile" or basename.startswith("Dockerfile."):
+        return True
+
+    if name.startswith("3rdparty"):
+        return True
+
+    if name.startswith(".txdev") or name.startswith(".claude"):
+        return True
+
+    if name in ALLOW_SPECIFIC_FILE:
+        return True
+
+    return False
+
+
+def main() -> None:
+    cmd = ["git", "ls-files"]
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    (out, _) = proc.communicate()
+    res = out.decode("utf-8")
+    assert proc.returncode == 0, f"{' '.join(cmd)} errored: {res}"
+    error_list = [f for f in res.split() if not filename_allowed(f)]
+
+    if error_list:
+        report = "------File type check report----\n"
+        report += "\n".join(error_list)
+        report += f"\nFound {len(error_list)} files that are not allowed\n"
+        report += "We do not check in binary files into the repo.\n"
+        report += (
+            "If necessary, please discuss with committers and "
+            "modify tests/lint/check_file_type.py to enable the file you need.\n"
+        )
+        sys.stderr.write(report)
+        sys.stderr.flush()
+        sys.exit(-1)
+
+    print("check_file_type.py: all checks passed..")
+
+
+if __name__ == "__main__":
+    main()

tests/nightly/python/test_nnapi/__init__.py

@@ -0,0 +1,17 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Infrastructure and tests for NNAPI"""

tests/nightly/python/test_nnapi/conftest.py

@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+
+import pytest
+
+from tvm import rpc
+
+
+def remote():
+    required_env = ("TVM_TRACKER_HOST", "TVM_TRACKER_PORT", "RPC_DEVICE_KEY")
+    missing = [name for name in required_env if name not in os.environ]
+    if missing:
+        pytest.skip(f"NNAPI remote environment unavailable: {', '.join(missing)} not set")
+    rpc_tracker_host = os.environ["TVM_TRACKER_HOST"]
+    rpc_tracker_port = int(os.environ["TVM_TRACKER_PORT"])
+    rpc_device_key = os.environ["RPC_DEVICE_KEY"]
+    tracker = rpc.connect_tracker(rpc_tracker_host, rpc_tracker_port)
+    remote = tracker.request(rpc_device_key, priority=0, session_timeout=600)
+    return remote, tracker

tests/nightly/python/test_nnapi/infrastructure.py

@@ -0,0 +1,134 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: RUF005
+
+import numpy as np
+
+import tvm
+import tvm.script.relax as R
+from tvm.relax.backend.contrib.nnapi import partition_for_nnapi
+from tvm.support import ndk, utils
+
+
+# pylint: disable=import-outside-toplevel,missing-function-docstring
+def reshape_matmul(mod: tvm.IRModule):
+    from tvm.relax import Expr
+    from tvm.relax.dpl import DFPattern, rewrite_call
+    from tvm.relax.dpl.pattern import is_op, wildcard
+
+    input0 = wildcard()
+    input1 = wildcard()
+    pattern = is_op("relax.matmul")(input0, input1)
+
+    def _rewriter(expr: Expr, matches: dict[DFPattern, Expr]):
+        i0 = matches[input0]
+        i1 = matches[input1]
+        if len(i0.ty.shape) == 2 and len(i1.ty.shape) == 2:
+            i0_shape = [1] + [*i0.ty.shape.values]
+            i1_shape = [1] + [*i1.ty.shape.values]
+            oshape = matches[pattern].ty.shape
+            return R.reshape(R.matmul(R.reshape(i0, i0_shape), R.reshape(i1, i1_shape)), oshape)
+        return expr
+
+    mod["main"] = rewrite_call(pattern, _rewriter, mod["main"])
+    return mod
+
+
+def decompose_clip(mod: tvm.IRModule) -> tvm.IRModule:
+    from tvm.relax import Expr
+    from tvm.relax.dpl import DFPattern, rewrite_call
+    from tvm.relax.dpl.pattern import is_op, wildcard
+
+    input_pattern = wildcard()
+    min_pattern = wildcard()
+    max_pattern = wildcard()
+    pattern = is_op("relax.clip")(input_pattern, min_pattern, max_pattern)
+
+    def _rewriter(expr: Expr, matches: dict[DFPattern, Expr]) -> Expr:  # pylint: disable=unused-argument
+        dtype = matches[input_pattern].ty.dtype
+        return R.minimum(
+            R.maximum(
+                matches[input_pattern],
+                R.const(np.array(matches[min_pattern].value.value).astype(dtype), dtype),
+            ),
+            R.const(np.array(matches[max_pattern].value.value).astype(dtype), dtype),
+        )
+
+    mod["main"] = rewrite_call(pattern, _rewriter, mod["main"])
+    return mod
+
+
+def _build(mod, enable_nnapi):
+    if isinstance(mod, tvm.ir.Call):
+        mod = tvm.IRModule.from_expr(mod)
+
+    if enable_nnapi:
+        mod = tvm.relax.transform.FoldConstant()(mod)
+        mod = reshape_matmul(mod)
+        mod = decompose_clip(mod)
+        mod = partition_for_nnapi(mod)
+
+        mod = tvm.relax.transform.RunCodegen()(mod)
+    ex = tvm.compile(mod, target={"kind": "llvm", "mtriple": "aarch64-linux-android"})
+
+    return ex
+
+
+def _run(remote, tracker, ex, inputs):
+    tmp = utils.tempdir()
+    so_name = "test_mod.so"
+    so_path = tmp / so_name
+    ex.export_library(str(so_path), fcompile=ndk.create_shared, options=["-shared", "-fPIC", "-lm"])
+
+    remote.upload(so_path)
+    dev = remote.cpu(0)
+
+    try:
+        # Execute the model on the remote.
+        remote_ex = remote.load_module(so_name)
+        vm = tvm.relax.VirtualMachine(remote_ex, device=dev)
+
+        inputs = [x.copyto(dev) for x in inputs]
+
+        vm.set_input("main", *inputs)
+        vm.invoke_stateful("main")
+        output = vm.get_outputs("main")
+        output = output.numpy()
+    except Exception as e:
+        # Re-raise all exceptions
+        raise e
+    finally:
+        # Manually close the connection.
+        # See https://discuss.tvm.apache.org/t/trouble-with-rpc-session/14008/.
+        #
+        # TODO: Remove if it does not happen on Python 3.11.
+        remote._sess.get_function("CloseRPCConnection")()
+        tracker.close()
+        pass
+
+    return output
+
+
+def build_and_run(
+    remote,
+    tracker,
+    mod,
+    inputs,
+    enable_nnapi=False,
+):
+    ex = _build(mod, enable_nnapi)
+    return _run(remote, tracker, ex, inputs)

tests/nightly/python/test_nnapi/test_network.py

@@ -0,0 +1,138 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""NNAPI network tests."""
+
+import numpy as np
+import pytest
+
+pytest.importorskip("onnx")
+
+import onnx
+
+import tvm
+from test_nnapi.conftest import remote
+from test_nnapi.infrastructure import build_and_run  # , build_and_run_vm
+from tvm.contrib.download import download_testdata
+from tvm.relax.frontend.onnx import from_onnx
+from tvm.testing import env
+
+
+def _build_and_run_network(remote_obj, tracker, mod, input_data):
+    """Helper function to build and run a network."""
+
+    def execute_on_host(mod, inputs):
+        with tvm.transform.PassContext(opt_level=3):
+            ex = tvm.compile(mod, target="llvm")
+        dev = tvm.cpu(0)
+        vm = tvm.relax.VirtualMachine(ex, device=dev)
+        output = vm["main"](*inputs)
+        return output.numpy()
+
+    outputs = []
+    for nnapi in [True, False]:
+        if nnapi:
+            outputs.append(
+                build_and_run(
+                    remote_obj,
+                    tracker,
+                    mod,
+                    input_data,
+                    enable_nnapi=nnapi,
+                )
+            )
+        else:
+            outputs.append(execute_on_host(mod, input_data))
+    return outputs
+
+
+def get_network(name, dtype, input_shape=(1, 3, 224, 224)):
+    def download_model(model_url, name):
+        model_path = download_testdata(model_url, name + ".onnx", module="onnx")
+        onnx_model = onnx.load(model_path)
+
+        shape_dict = {"x": input_shape}
+        mod = from_onnx(onnx_model, shape_dict)
+        return mod
+
+    def create_model(name):
+        if "vgg11" == name:
+            model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/vgg11_Opset18_timm/vgg11_Opset18.onnx"
+        elif "mobilenetv3" == name:
+            model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/mobilenetv3_large_100_miil_Opset17_timm/mobilenetv3_large_100_miil_Opset17.onnx"
+        elif "alexnet" == name:
+            model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/alexnet_Opset17_torch_hub/alexnet_Opset17.onnx"
+        elif "resnet50" == name:
+            model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/resnet50_Opset18_timm/resnet50_Opset18.onnx"
+        elif "resnet34" == name:
+            model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/resnet34_Opset18_timm/resnet34_Opset18.onnx"
+        elif "resnet18" == name:
+            model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/resnet18_Opset18_timm/resnet18_Opset18.onnx"
+        elif "squeezenet" == name:
+            model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/squeezenet1_1_Opset18_torch_hub/squeezenet1_1_Opset18.onnx"
+        elif "vgg16" == name:
+            model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/vgg16_Opset18_timm/vgg16_Opset18.onnx"
+        elif "vgg19" == name:
+            model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/vgg19_Opset18_timm/vgg19_Opset18.onnx"
+        else:
+            assert False, f"Not supported model {name}"
+
+        return download_model(model_url, name)
+
+    mod = create_model(name)
+    return mod, {"data": (input_shape, dtype)}
+
+
+@pytest.mark.parametrize(
+    "name",
+    [
+        "alexnet",
+        "vgg11",
+        "vgg16",
+        "vgg19",
+        "resnet18",
+        "resnet34",
+        "resnet50",
+        "squeezenet",
+        "mobilenetv3",
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        "float32",
+    ],
+)
+@pytest.mark.skipif(not env.build_flag_enabled("USE_NNAPI_CODEGEN"), reason="need nnapi")
+def test_network(name, dtype):
+    remote_obj, tracker = remote()
+    print(f"Network evaluating {name} with dtype {dtype}")
+    np.random.seed(0)
+    mod, inputs = get_network(name, dtype)
+    input_data = {}
+
+    for _name, (shape, _dtype) in inputs.items():
+        input_data[_name] = np.random.uniform(-1.0, 1.0, shape).astype(_dtype)
+
+    inputs_tvm: list[tvm.runtime.Tensor] = [tvm.runtime.tensor(v) for k, v in input_data.items()]
+    outputs = _build_and_run_network(remote_obj, tracker, mod, inputs_tvm)
+    nnapi_out = outputs[0]
+    expected_out = outputs[1]
+    tvm.testing.assert_allclose(nnapi_out, expected_out, rtol=1e-4, atol=1e-5)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/nightly/python/test_nnapi/test_ops.py

@@ -0,0 +1,361 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+"""NNAPI integration operator tests."""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.script
+import tvm.script.relax as R
+import tvm.script.tirx as T
+from test_nnapi.conftest import remote
+from test_nnapi.infrastructure import build_and_run
+
+
+def _build_and_run_network(remote_obj, tracker, mod, input_data):
+    """Helper function to build and run a network."""
+
+    def execute_on_host(mod, inputs):
+        with tvm.transform.PassContext(opt_level=3):
+            ex = tvm.compile(mod, target="llvm")
+        dev = tvm.cpu(0)
+        vm = tvm.relax.VirtualMachine(ex, device=dev)
+        output = vm["main"](*inputs)
+        return output.numpy()
+
+    outputs = []
+    for nnapi in [True, False]:
+        if nnapi:
+            outputs.append(
+                build_and_run(
+                    remote_obj,
+                    tracker,
+                    mod,
+                    input_data,
+                    enable_nnapi=nnapi,
+                )
+            )
+        else:
+            outputs.append(execute_on_host(mod, input_data))
+    return outputs
+
+
+@pytest.mark.parametrize(
+    "op",
+    [
+        R.exp,
+        R.log,
+        R.negative,
+        R.sqrt,
+        R.rsqrt,
+        R.floor,
+        R.nn.relu,
+        R.nn.softmax,
+        R.sigmoid,
+        R.tanh,
+        R.abs,
+    ],
+)
+def test_unary(op, input_shape=(1, 2, 8, 5)):
+    remote_obj, tracker = remote()
+
+    def create_model() -> tvm.IRModule:
+        @tvm.script.ir_module
+        class Module:
+            @R.function
+            def main(i0: R.Tensor((1, 2, 8, 5), "float32")) -> R.Tensor((1, 2, 8, 5), "float32"):
+                with R.dataflow():
+                    t0 = op(i0)
+                    R.output(t0)
+                return t0
+
+        return Module
+
+    mod = create_model()
+    verify(
+        remote_obj,
+        tracker,
+        mod,
+        inputs=[np.random.uniform(size=(1, 2, 8, 5)).astype("float32")],
+    )
+
+
+@pytest.mark.parametrize(
+    "op",
+    [
+        R.power,
+        R.greater,
+        R.add,
+        R.multiply,
+        R.subtract,
+        R.equal,
+        R.less,
+        R.less_equal,
+        R.not_equal,
+        R.maximum,
+        R.minimum,
+        R.greater_equal,
+    ],
+)
+def test_elementwise_binary(op, input_shape=(1, 2, 8, 5)):
+    remote_obj, tracker = remote()
+
+    def create_model() -> tvm.IRModule:
+        @tvm.script.ir_module
+        class Module:
+            @R.function
+            def main(
+                i0: R.Tensor((1, 2, 8, 5), "float32"),
+                i1: R.Tensor((1, 2, 8, 5), "float32"),
+            ) -> R.Tensor((1, 2, 8, 5), "float32"):
+                with R.dataflow():
+                    t0 = op(i0, i1)
+                    R.output(t0)
+                return t0
+
+        return Module
+
+    mod = create_model()
+    verify(
+        remote_obj,
+        tracker,
+        mod,
+        inputs=[
+            np.random.uniform(size=input_shape).astype("float32"),
+            np.random.uniform(size=input_shape).astype("float32"),
+        ],
+    )
+
+
+def test_divide(input_shape=(1, 2, 8, 5)):
+    remote_obj, tracker = remote()
+
+    def create_model(input_shape) -> tvm.IRModule:
+        @tvm.script.ir_module
+        class Module:
+            @R.function
+            def main(
+                i0: R.Tensor((1, 2, 8, 5), "float32"),
+                i1: R.Tensor((1, 2, 8, 5), "float32"),
+            ) -> R.Tensor((1, 2, 8, 5), "float32"):
+                with R.dataflow():
+                    t0 = R.divide(i0, i1)
+                    R.output(t0)
+                return t0
+
+        return Module
+
+    mod = create_model(input_shape)
+    verify(
+        remote_obj,
+        tracker,
+        mod,
+        inputs=[
+            np.random.uniform(size=input_shape).astype("float32"),
+            np.random.uniform(size=input_shape).astype("float32") + np.ones(input_shape, "float32"),
+        ],
+    )
+
+
+def test_matmul():
+    remote_obj, tracker = remote()
+
+    def create_model() -> tvm.IRModule:
+        @tvm.script.ir_module
+        class Module:
+            @R.function
+            def main(
+                i0: R.Tensor((5, 3, 4), "float32"),
+                i1: R.Tensor((5, 4, 8), "float32"),
+            ) -> R.Tensor((5, 3, 8), "float32"):
+                with R.dataflow():
+                    t0 = R.matmul(i0, i1)
+                    R.output(t0)
+                return t0
+
+        return Module
+
+    mod = create_model()
+    verify(
+        remote_obj,
+        tracker,
+        mod,
+        inputs=[
+            np.random.random(size=(5, 3, 4)).astype("float32"),
+            np.random.random(size=(5, 4, 8)).astype("float32"),
+        ],
+    )
+
+
+def test_permute_dims():
+    remote_obj, tracker = remote()
+
+    def create_model() -> tvm.IRModule:
+        @tvm.script.ir_module
+        class Module:
+            @R.function
+            def main(
+                i0: R.Tensor((5, 4, 8), "float32"),
+            ) -> R.Tensor((8, 5, 4), "float32"):
+                with R.dataflow():
+                    t0 = R.permute_dims(i0, axes=[2, 0, 1])
+                    R.output(t0)
+                return t0
+
+        return Module
+
+    mod = create_model()
+    verify(
+        remote_obj,
+        tracker,
+        mod,
+        inputs=[
+            np.random.random(size=(5, 4, 8)).astype("float32"),
+        ],
+    )
+
+
+def test_astype():
+    remote_obj, tracker = remote()
+
+    def create_model() -> tvm.IRModule:
+        @tvm.script.ir_module
+        class Module:
+            @R.function
+            def main(
+                i0: R.Tensor((8, 10, 15), "float32"),
+            ) -> R.Tensor((8, 10, 15), "float16"):
+                with R.dataflow():
+                    t0: R.Tensor((8, 10, 15), "float16") = R.astype(i0, dtype="float16")
+                    R.output(t0)
+                return t0
+
+        return Module
+
+    mod = create_model()
+    verify(
+        remote_obj,
+        tracker,
+        mod,
+        inputs=[
+            tvm.runtime.tensor(np.random.uniform(size=(8, 10, 15)).astype("float32")),
+        ],
+    )
+
+
+def test_mean():
+    remote_obj, tracker = remote()
+
+    def create_model() -> tvm.IRModule:
+        @tvm.script.ir_module
+        class Module:
+            @R.function
+            def main(
+                i0: R.Tensor((1, 10, 15), "float32"),
+            ) -> R.Tensor((1, 10, 1), "float32"):
+                n = T.int64()
+                with R.dataflow():
+                    t0: R.Tensor((1, 10, 1), "float32") = R.mean(i0, axis=[-1], keepdims=True)
+                    R.output(t0)
+                return t0
+
+        return Module
+
+    mod = create_model()
+    verify(
+        remote_obj,
+        tracker,
+        mod,
+        inputs=[
+            tvm.runtime.tensor(np.random.uniform(size=(1, 10, 15)).astype("float32")),
+        ],
+    )
+
+
+def test_conv2d():
+    remote_obj, tracker = remote()
+
+    def create_model() -> tvm.IRModule:
+        @tvm.script.ir_module
+        class Module:
+            @R.function
+            def main(
+                i0: R.Tensor((1, 3, 224, 224), "float32"),
+                i1: R.Tensor((64, 3, 3, 3), "float32"),
+                i2: R.Tensor((1, 64, 1, 1), "float32"),
+            ):
+                with R.dataflow():
+                    t0 = R.nn.conv2d(i0, i1, strides=(1, 1), padding=(1, 1))
+                    t0 = R.add(i2, t0)
+                    R.output(t0)
+                return t0
+
+        return Module
+
+    mod = create_model()
+    verify(
+        remote_obj,
+        tracker,
+        mod,
+        inputs=[
+            np.random.random(size=(1, 3, 224, 224)).astype("float32"),
+            np.random.random(size=(64, 3, 3, 3)).astype("float32"),
+            np.random.random(size=(1, 64, 1, 1)).astype("float32"),
+        ],
+    )
+
+
+def test_max_pool2d():
+    remote_obj, tracker = remote()
+
+    def create_model() -> tvm.IRModule:
+        @tvm.script.ir_module
+        class Module:
+            @R.function
+            def main(
+                i0: R.Tensor((1, 1, 28, 28), "float32"),
+            ):
+                with R.dataflow():
+                    t0 = R.nn.max_pool2d(i0, pool_size=(1, 1), strides=(1, 1), padding=(0, 0))
+                    R.output(t0)
+                return t0
+
+        return Module
+
+    mod = create_model()
+    verify(
+        remote_obj,
+        tracker,
+        mod,
+        inputs=[
+            np.random.random(size=(1, 1, 28, 28)).astype("float32"),
+        ],
+    )
+
+
+def verify(remote_obj, tracker, mod, inputs):
+    inputs_tvm: list[tvm.runtime.Tensor] = [tvm.runtime.tensor(v) for v in inputs]
+    outputs = _build_and_run_network(remote_obj, tracker, mod, inputs_tvm)
+    nnapi_out = outputs[0]
+    expected_out = outputs[1]
+    tvm.testing.assert_allclose(nnapi_out, expected_out, rtol=1e-4, atol=1e-5)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/all-platform-minimal-test/README.md

@@ -0,0 +1,29 @@
+<!--- Licensed to the Apache Software Foundation (ASF) under one -->
+<!--- or more contributor license agreements.  See the NOTICE file -->
+<!--- distributed with this work for additional information -->
+<!--- regarding copyright ownership.  The ASF licenses this file -->
+<!--- to you under the Apache License, Version 2.0 (the -->
+<!--- "License"); you may not use this file except in compliance -->
+<!--- with the License.  You may obtain a copy of the License at -->
+
+<!---   http://www.apache.org/licenses/LICENSE-2.0 -->
+
+<!--- Unless required by applicable law or agreed to in writing, -->
+<!--- software distributed under the License is distributed on an -->
+<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
+<!--- KIND, either express or implied.  See the License for the -->
+<!--- specific language governing permissions and limitations -->
+<!--- under the License. -->
+
+# Core Cross Platform Regression Tests
+
+CI Unit test cases that will run on all platforms.
+To reduce the CI burden, we only put in test-cases that are platform sensitive.
+Please use the following guideline:
+
+- Always consider add tests to the unittest folder first.
+- If a problems that passes the Linux pipeline but fails in Windows or MacOS,
+  we should isolate the problem, write a minimal regression test case
+  and add it to this folder.
+- A test case in this folder should be minimal and finish in a reasonable amount of time.
+- Document about why it should be in the all-platform-minimal-test.

tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py

@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+"""LLVM enablement tests."""
+
+import ctypes
+import math
+import re
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te, topi
+from tvm.support import utils
+from tvm.testing import env
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_llvm_add_pipeline():
+    """all-platform-minimal-test: Check LLVM enablement."""
+    nn = 128
+    n = tvm.runtime.convert(nn)
+    A = te.placeholder((n,), name="A")
+    B = te.placeholder((n,), name="B")
+    AA = te.compute((n,), lambda *i: A(*i), name="A")
+    BB = te.compute((n,), lambda *i: B(*i), name="B")
+    T = te.compute(A.shape, lambda *i: AA(*i) + BB(*i), name="T")
+    C = te.compute(A.shape, lambda *i: T(*i), name="C")
+
+    sch = tvm.s_tir.Schedule(te.create_prim_func([A, B, C]))
+    xo, xi = sch.split(sch.get_loops("C")[0], factors=[None, 4])
+    sch.parallel(xo)
+    sch.vectorize(xi)
+
+    def check_llvm():
+        # BUILD and invoke the kernel.
+        f = tvm.compile(sch.mod, target="llvm")
+        dev = tvm.cpu(0)
+        # launch the kernel.
+        n = nn
+        a = tvm.runtime.tensor(np.random.uniform(size=n).astype(A.dtype), dev)
+        b = tvm.runtime.tensor(np.random.uniform(size=n).astype(B.dtype), dev)
+        c = tvm.runtime.tensor(np.zeros(n, dtype=C.dtype), dev)
+        f(a, b, c)
+        tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
+
+    check_llvm()

tests/python/all-platform-minimal-test/test_validate_runtime_library.py

@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Post-install checks for a built TVM wheel.
+
+These live in ``tests/python/all-platform-minimal-test`` so the standard suite and
+the cibuildwheel ``test-command`` run a single pytest invocation. The assertions
+here are wheel-specific things the rest of the suite cannot check -- that LLVM is
+enabled (the other LLVM test merely *skips* when LLVM is absent) and that the CUDA
+runtime library got bundled -- so each is gated behind a ``TVM_WHEEL_EXPECT_*`` env
+var and SKIPS unless that var is set. cibuildwheel sets the vars (see
+``CIBW_TEST_ENVIRONMENT`` in ``.github/actions/build-wheel-for-publish``); ordinary
+source-build CI (e.g. ``main.yml``) leaves them unset, so these tests skip there and
+never fail a non-wheel / non-LLVM / non-CUDA build.
+"""
+
+import glob
+import os
+from pathlib import Path
+
+import pytest
+
+import tvm
+
+
+def test_llvm_enabled():
+    """Every published TVM wheel ships with LLVM enabled. Only assert this when
+    validating a wheel (``TVM_WHEEL_EXPECT_LLVM=1``); skip otherwise so source
+    builds with LLVM off do not fail."""
+    if os.environ.get("TVM_WHEEL_EXPECT_LLVM") != "1":
+        pytest.skip("LLVM enablement only asserted during wheel validation")
+    assert tvm.runtime.enabled("llvm"), "wheel was not built with LLVM enabled"
+
+
+def test_cuda_runtime_present():
+    """The bundled CUDA runtime library must be present in tvm/lib."""
+    if os.environ.get("TVM_WHEEL_EXPECT_CUDA_RUNTIME") != "1":
+        pytest.skip("CUDA runtime not expected in this wheel")
+    libdir = Path(tvm.__file__).resolve().parent / "lib"
+    present = glob.glob(str(libdir / "libtvm_runtime_cuda.*")) or glob.glob(
+        str(libdir / "tvm_runtime_cuda.*")
+    )
+    assert present, "CUDA runtime expected but not bundled in tvm/lib"

tests/python/arith/test_arith_analyzer_object.py

@@ -0,0 +1,287 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.arith.analyzer import CompareResult, Extension
+from tvm.runtime import Object
+
+
+def test_analyzer_is_ffi_object_with_persistent_state():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+
+    assert isinstance(analyzer, Object)
+
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+    assert analyzer.const_int_bound_is_bound(x)
+    assert analyzer.can_prove(x < 8)
+    assert not analyzer.can_prove(x < 4)
+
+    bound = analyzer.const_int_bound(x + 1)
+    assert bound.min_value == 1
+    assert bound.max_value == 8
+
+
+def test_analyzer_object_constraint_scope_and_override_bind():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+
+    with analyzer.constraint_scope(x % 3 == 0):
+        assert analyzer.modular_set(x).coeff == 3
+
+    assert analyzer.modular_set(x).coeff != 3
+
+    analyzer = tvm.arith.Analyzer()
+    y = tirx.Var("y", "int64")
+    analyzer.bind(y, tirx.const(4, "int64"))
+    tvm.ir.assert_structural_equal(analyzer.simplify(y + 1), tirx.const(5, "int64"))
+
+    analyzer.bind(y, tirx.const(8, "int64"), allow_override=True)
+    tvm.ir.assert_structural_equal(analyzer.simplify(y + 1), tirx.const(9, "int64"))
+
+
+def test_analyzer_object_update_const_int_bound():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+
+    analyzer.update(x, tvm.arith.ConstIntBound(2, 5))
+
+    bound = analyzer.const_int_bound(x + 1)
+    assert bound.min_value == 3
+    assert bound.max_value == 6
+
+
+def test_analyzer_object_update_modular_set():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+
+    assert analyzer.modular_set(x).coeff == 1
+    analyzer.update(x, tvm.arith.ModularSet(4, 0))
+
+    result = analyzer.modular_set(x)
+    assert result.coeff == 4
+    assert result.base == 0
+
+
+def test_analyzer_object_update_int_set():
+    analyzer = tvm.arith.Analyzer()
+    y = tirx.Var("y", "int32")
+
+    analyzer.update(y, tvm.arith.IntervalSet(0, 8))
+
+    int_set = analyzer.int_set(y)
+    assert int_set.min_value.value == 0
+    assert int_set.max_value.value == 8
+
+
+def test_analyzer_object_update_rejects_unknown_info():
+    analyzer = tvm.arith.Analyzer()
+    y = tirx.Var("y", "int32")
+
+    with pytest.raises(TypeError):
+        analyzer.update(y, "not-an-info-object")
+
+
+def test_analyzer_object_can_prove_comparison_predicates():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+
+    assert analyzer.can_prove(x >= 0)
+    assert not analyzer.can_prove(x >= 1)
+    assert analyzer.can_prove(x < 8)
+    assert not analyzer.can_prove(x < 7)
+
+
+def test_analyzer_object_update_const_int_bound_half_space():
+    analyzer = tvm.arith.Analyzer()
+    n = tirx.Var("n", "int32")
+
+    assert not analyzer.can_prove(n >= 0)
+    analyzer.update(n, tvm.arith.ConstIntBound(0, tvm.arith.ConstIntBound.POS_INF))
+    assert analyzer.can_prove(n >= 0)
+
+
+def test_analyzer_object_int_set_from_bound_vars():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+
+    int_set = analyzer.int_set(x + 1)
+    assert int_set.min_value.value == 1
+    assert int_set.max_value.value == 8
+
+
+def test_analyzer_object_set_maximum_rewrite_steps():
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    expr = (x + y) * 2 - x * 2 - y * 2 + tirx.max(x, y) - tirx.min(x, y)
+
+    capped = tvm.arith.Analyzer()
+    capped.set_maximum_rewrite_steps(1)
+    with pytest.raises(RuntimeError):
+        capped.rewrite_simplify(expr)
+
+    # A generous limit must not interfere with normal simplification.
+    relaxed = tvm.arith.Analyzer()
+    relaxed.set_maximum_rewrite_steps(1000)
+    relaxed.rewrite_simplify(expr)
+
+
+def test_analyzer_object_try_compare_transitive():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    z = tirx.Var("z", "int32")
+
+    assert analyzer.try_compare(x, y) == CompareResult.UNKNOWN
+
+    with analyzer.constraint_scope(x < y):
+        with analyzer.constraint_scope(y < z):
+            # Direct known comparison.
+            assert analyzer.try_compare(x, y) == CompareResult.LT
+            # Transitive chain x < y < z is found only when propagation is enabled.
+            assert analyzer.try_compare(x, z) == CompareResult.LT
+            assert analyzer.try_compare(x, z, propagate_inequalities=False) == CompareResult.UNKNOWN
+
+
+def test_analyzer_object_enabled_extensions_round_trip():
+    analyzer = tvm.arith.Analyzer()
+
+    assert analyzer.enabled_extensions == Extension.NoExtensions
+
+    analyzer.enabled_extensions = Extension.ComparisonOfProductAndSum
+    assert analyzer.enabled_extensions == Extension.ComparisonOfProductAndSum
+
+    analyzer.enabled_extensions = Extension.NoExtensions
+    assert analyzer.enabled_extensions == Extension.NoExtensions
+
+
+def test_analyzer_object_rewrite_simplify_stats():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int32")
+
+    analyzer.reset_rewrite_simplify_stats()
+    assert analyzer.rewrite_simplify_stats.nodes_visited == 0
+
+    analyzer.rewrite_simplify(x + 0)
+    assert analyzer.rewrite_simplify_stats.nodes_visited > 0
+
+    analyzer.reset_rewrite_simplify_stats()
+    assert analyzer.rewrite_simplify_stats.nodes_visited == 0
+
+
+def test_analyzer_object_state_persists_across_ffi_calls():
+    analyzer = tvm.arith.Analyzer()
+    tile = tirx.Var("tile", "int32")
+    i = tirx.Var("i", "int32")
+    analyzer.bind(tile, tvm.tirx.const(8, "int32"))
+
+    # The same analyzer object is borrowed by the C++ DetectIterMap entry point;
+    # its binding makes the otherwise-undetectable floormod recognizable.
+    result = tvm.arith.detect_iter_map([i % tile], {i: tvm.ir.Range(0, 32)}, analyzer=analyzer)
+    assert len(result.indices) == 1
+
+    # The binding still lives in the same stateful object after the FFI call.
+    tvm.ir.assert_structural_equal(analyzer.simplify(tile), tvm.tirx.const(8, "int32"))
+
+
+def test_analyzer_object_clone_is_independent():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+    y = tirx.Var("y", "int64")
+    z = tirx.Var("z", "int64")
+
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+
+    clone = analyzer.clone()
+    assert clone is not analyzer
+    assert clone.can_prove(x < 8)
+
+    clone.bind(y, tvm.ir.Range(0, 4))
+    assert clone.can_prove(y < 4)
+    assert not analyzer.can_prove(y < 4)
+
+    analyzer.bind(z, tvm.ir.Range(0, 4))
+    assert analyzer.can_prove(z < 4)
+    assert not clone.can_prove(z < 4)
+
+    assert analyzer.can_prove(x < 8)
+    assert clone.can_prove(x < 8)
+
+
+def test_analyzer_object_clone_copies_every_sub_analyzer():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+    w = tirx.Var("w", "int64")
+    v = tirx.Var("v", "int64")
+
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+    analyzer.update(x, tvm.arith.ModularSet(4, 0))
+    analyzer.bind(w, tirx.const(4, "int64"))
+    analyzer.update(v, tvm.arith.IntervalSet(2, 9))
+    analyzer.enabled_extensions = Extension.ComparisonOfProductAndSum
+
+    clone = analyzer.clone()
+
+    assert clone.can_prove(x < 8)
+    assert clone.modular_set(x).coeff == 4
+    tvm.ir.assert_structural_equal(clone.simplify(w + 1), tirx.const(5, "int64"))
+    assert clone.int_set(v).max_value.value == 9
+    assert clone.enabled_extensions == Extension.ComparisonOfProductAndSum
+    assert clone.try_compare(x, tirx.const(0, "int64")) == CompareResult.GE
+
+    t = tirx.Var("t", "int64")
+    clone.update(x, tvm.arith.ModularSet(8, 0), override=True)
+    clone.update(v, tvm.arith.IntervalSet(0, 3), override=True)
+    clone.bind(w, tirx.const(8, "int64"), allow_override=True)
+    clone.bind(t, tvm.ir.Range(0, 4))
+    clone.enabled_extensions = Extension.NoExtensions
+
+    assert analyzer.modular_set(x).coeff == 4
+    assert clone.modular_set(x).coeff == 8
+    assert analyzer.int_set(v).max_value.value == 9
+    assert clone.int_set(v).max_value.value == 3
+    tvm.ir.assert_structural_equal(analyzer.simplify(w + 1), tirx.const(5, "int64"))
+    tvm.ir.assert_structural_equal(clone.simplify(w + 1), tirx.const(9, "int64"))
+    assert analyzer.enabled_extensions == Extension.ComparisonOfProductAndSum
+    assert clone.enabled_extensions == Extension.NoExtensions
+    assert clone.try_compare(t, tirx.const(0, "int64")) == CompareResult.GE
+    assert analyzer.try_compare(t, tirx.const(0, "int64")) == CompareResult.UNKNOWN
+
+
+def test_analyzer_object_clone_resets_rewrite_stats():
+    analyzer = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+    y = tirx.Var("y", "int64")
+    analyzer.bind(x, tvm.ir.Range(0, 8))
+    analyzer.bind(y, tvm.ir.Range(0, 8))
+    analyzer.simplify((x + y) * 2 - x - y)
+    source_attempts = analyzer.rewrite_simplify_stats.rewrites_attempted
+    assert source_attempts > 0
+
+    clone = analyzer.clone()
+    assert clone.rewrite_simplify_stats.rewrites_attempted == 0
+    assert analyzer.rewrite_simplify_stats.rewrites_attempted == source_attempts
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_canonical_simplify.py

@@ -0,0 +1,538 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E731, F841
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import te, tirx
+from tvm.script import tirx as T
+
+
+class CanonicalChecker:
+    def __init__(self):
+        self.analyzer = tvm.arith.Analyzer()
+
+    def _convert(self, expr):
+        # TODO(Lunderberg): Make utility functions `tirx.convert` and
+        # `relax.convert` that convert to their respective IR types.
+        # Implementation should be in C++, and should only consist of
+        # conversions that are applied automatically through FFI.
+        if isinstance(expr, int):
+            return T.int32(expr)
+        else:
+            return expr
+
+    def verify(self, data, expected):
+        res = self.analyzer.canonical_simplify(data)
+        expected = self._convert(expected)
+        assert tvm_ffi.structural_equal(res, expected), (
+            f"\ndata={data}\nres={res}\nexpected={expected}"
+        )
+
+
+def test_mul_sum_simplify():
+    ck = CanonicalChecker()
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+
+    ck.verify(2 + (3 * x + z + y + 1) * 4 + x, x * 13 + z * 4 + y * 4 + 6)
+    ck.verify(x * 3 - 4 * x + 1, 1 - x)
+    ck.verify(y + x * 3 - 5 * x + 1 + y, y * 2 + 1 - x * 2)
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+    # trucdiv
+    ck.verify(tdiv(x + y + x + y * 3, 2), y * 2 + x)
+    ck.verify(tmod(x + y + x + y * 3, 2), 0)
+
+    # floordiv
+    fld = tvm.tirx.floordiv
+    flm = tvm.tirx.floormod
+    ck.verify(flm(x + x + y * 3, 2), flm(y * 3, 2))
+    ck.verify(fld(x + y + x + y * 3, 2), y * 2 + x)
+    ck.verify(flm(x + y + x + y * 3, 2), 0)
+    ck.verify(fld(x + x + y * 3, 2), fld(y * 3, 2) + x)
+
+
+def test_split_index_simplify():
+    ck = CanonicalChecker()
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+
+    # trucdiv
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+
+    # split div const
+    ck.verify(tdiv(x, 3) * 3 + tmod(x, 3), x)
+    ck.verify(tdiv(x, 6) * 6 + tmod(tdiv(x, 3), 2) * 3 + tmod(x, 3), x)
+    ck.verify(tdiv(tdiv(tmod(x, 16), 2) * 2, 4), tdiv(tmod(x, 16), 4))
+    ck.verify(tdiv(tmod(x, 2), 8), 0)
+    ck.verify(tdiv(tmod(x, 2), 7), 0)
+    ck.verify(tdiv(tdiv(tmod(x, 16), 2) * 2, 6), tdiv(tmod(x, 16), 6))
+
+    # split mod const
+    ck.verify(tmod((x * 8), 16), tmod(x, 2) * 8)
+    ck.verify(tmod(x * 8, 2), 0)
+
+    # simplify then fold
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 1000))
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 1000))
+    ck.verify(tdiv(x * 4 + y, 2) * 2 + tmod(x * 4 + y, 2), x * 4 + y)
+    # complex fold
+    ck.verify(tdiv(z * 9 + y, 2) * 2 + tmod(z * 9 + y, 2), z * 9 + y)
+
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(-100, 1000), True)
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(-100, 1000), True)
+    ck.verify(tdiv(x * 4 + y, 2) * 2 + tmod(x * 4 + y, 2), x * 4 + y)
+
+    # floordiv
+    fld = tvm.tirx.floordiv
+    flm = tvm.tirx.floormod
+    ck.verify(fld(x * 5, 2), fld(x * 5, 2))
+    ck.verify(fld(x, 3) * 3 + flm(x, 3), x)
+    ck.verify(fld(x, 6) * 6 + flm(fld(x, 3), 2) * 3 + flm(x, 3), x)
+    ck.verify(fld(fld(flm(x, 16), 2) * 2, 4), fld(flm(x, 16), 4))
+    ck.verify(fld(flm(x, 2), 8), 0)
+    ck.verify(fld(flm(x, 2), 7), 0)
+    ck.verify(fld(fld(flm(x, 16), 2) * 2, 6), fld(flm(x, 16), 6))
+
+    # floordiv(floormod(sum, m*n), n) => floormod(floordiv(sum, n), m)
+    # when sum has parts divisible by n
+    d_tile = te.var("d_tile")
+    i = te.var("i")
+    v = te.var("v")
+    ck.analyzer.update(d_tile, tvm.arith.ConstIntBound(0, 7), True)
+    ck.analyzer.update(i, tvm.arith.ConstIntBound(0, 1), True)
+    ck.analyzer.update(v, tvm.arith.ConstIntBound(0, 7), True)
+    ck.verify(fld(flm(d_tile * 16 + i * 8 + v, 64), 8), flm(d_tile * 2 + i, 8))
+
+    # cannot simplify mixed case, unless we canonicalize into one mode.
+    ck.verify(tdiv(x, 6) * 2 + tmod(fld(x, 3), 2), tdiv(x, 6) * 2 + tmod(fld(x, 3), 2))
+
+    ck.verify(tmod(-x, 2), tmod(x, -2) * -1)
+
+
+def test_div_simplify():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+    tdiv = tvm.tirx.truncdiv
+
+    # truc div
+    ck.verify(tdiv(16 + 48 * x, 16), x * 3 + 1)
+    # (17+48*x)/16 is not simplifiable for arbitrary x because when 17+48*x<0
+    # (17+48*x)/16 != 1+3*x
+    ck.verify(tdiv(17 + 48 * x, 16), tdiv(x * 48 + 17, 16))
+    # However, when x >= 0, then 17+48*x >= 0 and (17+48*x)/16 can be simplified
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 10))
+    ck.verify(tdiv(17 + 48 * x, 16), x * 3 + 1)
+    # Trying expressions that are not simplifiable for any values of the variables
+    ck.verify(tdiv(17 + 47 * x, 16), tdiv(x * 47 + 17, 16))
+
+    # floordiv
+    fld = tvm.tirx.floordiv
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(-1000, 10000), True)
+    ck.verify(fld(16 + 48 * x, 16), x * 3 + 1)
+    ck.verify(fld(17 + 48 * x, 16), x * 3 + 1)
+    ck.verify(fld(17 + 47 * x, 16), fld(x * 47 + 17, 16))
+
+
+def test_fp16_const_fold():
+    ck = CanonicalChecker()
+    zero = tvm.tirx.const(0, "float16")
+    one = tvm.tirx.const(1, "float16")
+    half = tvm.tirx.const(0.5, "float16")
+
+    ck.verify(zero + half, half)
+    ck.verify(half - zero, half)
+
+    ck.verify(zero * half, zero)
+    ck.verify(half * one, half)
+
+    ck.verify(half / one, half)
+    ck.verify(zero / half, zero)
+
+
+def test_floormod_simplify():
+    ck = CanonicalChecker()
+    flm = tvm.tirx.floormod
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(flm(flm((x * 4) + y - 466036, 24528) - 24512, 16), flm((x * 4) + y + 12, 16))
+    ck.verify(flm(flm((x * 4), 16), 8), flm(x, 2) * 4)
+
+    ck.verify(flm(-x, 2), flm(x, -2) * -1)
+
+
+def test_canonical_mixed():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+    z = tvm.tirx.const(3, "int32")
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+    ck.verify(tdiv(x, (z * z)) - tdiv(x, (z * z)), 0)
+    ck.verify(tdiv(x, (z + z)) - tdiv(x, (z + z)), 0)
+    ck.verify(x - 2 < 3, x < 5)
+    ck.verify(tvm.tirx.max(x, 1) - tvm.tirx.max(x, 1), 0)
+    ck.verify(tvm.tirx.min(x, 1) - tvm.tirx.min(x, 1), 0)
+    ck.verify(x * x - x * x, 0)
+    ck.verify(tmod(tdiv(tmod(x, 20), 2) * 2, 4), tdiv(tmod(x, 4), 2) * 2)
+
+    fld = tvm.tirx.floordiv
+    ck.verify(fld(x, (z * z)) - fld(x, (z * z)), 0)
+    ck.verify(fld(x, (z + z)) - fld(x, (z + z)), 0)
+
+
+def test_reduce_combiner_simplify():
+    ck = CanonicalChecker()
+    dummy = tvm.tirx.Var("dummy", "int32")
+    comm_reducer = te.comm_reducer
+    prod = comm_reducer(lambda x, y: x * y, lambda t0: tvm.tirx.const(1, t0))
+
+    sum_or_prod = comm_reducer(
+        lambda x, y: tvm.tirx.Select(dummy < 0, x + y, x * y),
+        lambda t0: tvm.tirx.Select(dummy < 0, tvm.tirx.const(0, t0), tvm.tirx.const(1, t0)),
+    )
+    sum_and_prod = comm_reducer(
+        lambda x, y: (x[0] + y[0], x[1] * y[1]),
+        lambda t0, t1: (tvm.tirx.const(0, t0), tvm.tirx.const(5, t1) - tvm.tirx.const(4, t1)),
+    )
+    some_reducer1 = comm_reducer(
+        lambda x, y: (
+            x[0] + y[0],
+            x[0] + y[0] + x[1] + y[1],
+            x[0] * y[2] + y[0] * x[2],
+            x[1] + y[2],
+            4.0,
+        ),
+        lambda t0, t1, t2, t3, t4: (
+            tvm.tirx.const(0, t0),
+            tvm.tirx.const(1, t1),
+            tvm.tirx.const(2, t2),
+            tvm.tirx.const(3, t3),
+            tvm.tirx.const(4, t4),
+        ),
+    )
+
+    k = te.reduce_axis((0, 10), name="k")
+    A = te.placeholder((10,), name="A")
+    # Test that SimplifyCombiner makes use of vranges
+    ck.analyzer.update(dummy, tvm.arith.ConstIntBound(-10, -4))
+    ck.verify(sum_or_prod(A[k], k), te.sum(A[k], k))
+    ck.verify(sum_or_prod(A[k], k, init=1), te.sum(A[k], k, init=1))
+    ck.analyzer.update(dummy, tvm.arith.ConstIntBound(5, 9), True)
+    ck.verify(sum_or_prod(A[k], k), prod(A[k], k))
+    ck.verify(sum_or_prod(A[k], k, init=1), prod(A[k], k, init=1))
+    ck.analyzer.update(dummy, tvm.arith.ConstIntBound(-10, 100), True)
+    ck.verify(sum_and_prod((A[k], A[10 - k]), k)[0], te.sum(A[k], k))
+    ck.verify(sum_and_prod((A[k], A[10 - k]), k)[1], prod(A[10 - k], k))
+
+    reference_simplified_sources = [
+        [A[0]],
+        [A[0], A[1]],
+        [A[0], A[2]],
+        [A[0], A[1], A[2], A[3]],
+        [A[4]],
+    ]
+    for j in range(5):
+        # Here we use the j-th component of the result, so only it and the components it
+        # depends on are left.
+        simplified = ck.analyzer.canonical_simplify(
+            some_reducer1((A[0], A[1], A[2], A[3], A[4]), k)[j]
+        )
+
+        # Check that the remaining components are the expected ones.
+        for lhs, rhs in zip(simplified.source, reference_simplified_sources[j]):
+            tvm.ir.assert_structural_equal(lhs, rhs)
+
+    # Test that components with side effects are not removed
+    dummy = tvm.ir.GlobalVar("dummy")
+    side_effect = lambda *xs: tvm.ir.Call(dummy, xs, ret_ty="int32")
+    ck.verify(
+        sum_and_prod((A[k], side_effect(A[10 - k])), k)[0],
+        sum_and_prod((A[k], side_effect(A[10 - k])), k)[0],
+    )
+    ck.verify(sum_and_prod((side_effect(A[k]), A[10 - k]), k)[0], te.sum(side_effect(A[k]), k))
+
+
+def test_reduce_simplify():
+    ck = CanonicalChecker()
+    k = te.reduce_axis((0, 10), name="k")
+    j = te.reduce_axis((-5, 3), name="j")
+    A = te.placeholder((10,), name="A")
+    ck.verify(te.sum(tvm.tirx.Select(k + j < 12, k + j, 0), [k, j]), te.sum(k + j, [k, j]))
+    ck.verify(te.sum(A[3], []), A[3])
+    ck.verify(te.sum(A[3], [], where=k > 12, init=1.0), tvm.tirx.const(1.0, dtype="float32"))
+    # The rule below is not typical, removed for now
+    ck.verify(te.sum(tvm.tirx.div(k, 10), k), te.sum(tvm.tirx.const(0, "int32"), k))
+
+
+def test_simplify_if_then_else():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+    # simplification that takes condition into account.
+    res = tvm.tirx.if_then_else(
+        (x * 4 + y) >= 466036,
+        tvm.tirx.if_then_else(
+            24512 <= tmod(((x * 4) + y) - 466036, 24528),
+            tmod(tmod(((x * 4) + y) - 466036, 24528) - 24512, 16),
+            x,
+        ),
+        y,
+    )
+
+    res2 = tvm.tirx.if_then_else(
+        (x * 4) >= 466036 - y,
+        tvm.tirx.if_then_else(
+            24512 <= tmod(((x * 4) + y) - 466036, 24528),
+            tmod(tmod(((x * 4) + y) - 466036, 24528) - 24512, 16),
+            x,
+        ),
+        y,
+    )
+    expected = tvm.tirx.if_then_else(
+        tvm.tirx.LE(466036, (x * 4 + y)),
+        tvm.tirx.if_then_else(
+            tvm.tirx.LE(24512, tmod(((x * 4) + y) - 4, 24528)), tmod(((x * 4) + y) - 4, 16), x
+        ),
+        y,
+    )
+    ck.verify(res, expected)
+    ck.verify(res2, expected)
+    # can only simplify if condition
+    res = tvm.tirx.Select(tvm.tirx.all(x >= -1, y >= 0), tmod(x + y + 100, 3), tmod(x + 100, 3))
+    expected = tvm.tirx.Select(tvm.tirx.all(x >= -1, y >= 0), tmod(x + y + 1, 3), tmod(x + 100, 3))
+    ck.verify(res, ck.analyzer.canonical_simplify(expected))
+
+    res = tvm.tirx.Select(x >= 10, tvm.tirx.if_then_else(tdiv(x, 3) > 2, x, 0), 0)
+    expected = tvm.tirx.Select(x >= 10, x, 0)
+    ck.verify(res, ck.analyzer.canonical_simplify(expected))
+
+    res = tvm.tirx.Select(x >= 10, tvm.tirx.if_then_else(tdiv(x, 3) < 2, x, 0), 0)
+    ck.verify(res, 0)
+
+
+def test_complex_cases():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+    res2 = (
+        tdiv(tdiv(tmod(x * 128 + y, 1296), 36) * 2 + 1, 2) * 36
+        + tdiv(tmod((x * 128) + y, 36) * 2 + 1, 2)
+        - tmod((x * 128) + y, 1296)
+        + 1
+    )
+    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 5))
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 127))
+    ck.verify(res2, 1)
+
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 1024), True)
+    res3 = (
+        tdiv(x * 1024 + y, 65536)
+        + tdiv(tmod(x * 1024 + y, 65536), 256)
+        + tdiv(tmod(x * 1024 + y, 256), 16)
+        + tmod(x * 1024 + y, 16)
+        - tdiv(y, 256)
+        - tdiv(tmod(y, 256), 16)
+        - tmod(y, 16)
+        - (x * 4)
+    )
+    ck.verify(res3, tdiv((x * 1024) + y, 256) - tdiv(y, 256) - (x * 4))
+
+
+def test_simplify_cast():
+    ck = CanonicalChecker()
+    tcast = tvm.tirx.Cast
+    fld = tvm.tirx.floordiv
+    flm = tvm.tirx.floormod
+    # cast(i64, i + j + 1) - cast(i64, i)
+    i = tvm.tirx.Var("i", "int32")
+    j = tvm.tirx.Var("j", "int32")
+    res = tcast("int64", i + j + 1) - tcast("int64", i)
+    ck.verify(res, tcast("int64", j) + tvm.tirx.const(1, "int64"))
+    # cast(i32, i + j + 1) - cast(i32, i)
+    i = tvm.tirx.Var("i", "int64")
+    j = tvm.tirx.Var("j", "int64")
+    ck.analyzer.update(i, tvm.arith.ConstIntBound(0, 10))
+    ck.analyzer.update(j, tvm.arith.ConstIntBound(0, 10))
+    res = tcast("int32", i + j + 1) - tcast("int32", i)
+    ck.verify(res, tcast("int32", j) + 1)
+    # cast(i32, i + j - 100)
+    i = tvm.tirx.Var("i", "int64")
+    j = tvm.tirx.Var("j", "int64")
+    ck.analyzer.update(i, tvm.arith.ConstIntBound(0, 2**31 - 1))
+    ck.analyzer.update(j, tvm.arith.ConstIntBound(0, 10))
+    res = tcast("int32", i + j - 100)
+    ck.verify(res, res)
+    # cast(i32, flm(axis, 7i64) * 2i64 + 1i64) + 1i32
+    # - cast(i32, flm(axis, 7i64) * 2i64)
+    axis = tvm.tirx.Var("axis", "int64")
+    ck.analyzer.update(axis, tvm.arith.ConstIntBound(0, 42))
+    res = (
+        tcast(
+            "int32",
+            flm(axis, tvm.tirx.const(7, "int64")) * tvm.tirx.const(2, "int64")
+            + tvm.tirx.const(1, "int64"),
+        )
+        + tvm.tirx.const(1, "int32")
+        - tcast("int32", flm(axis, tvm.tirx.const(7, "int64")) * tvm.tirx.const(2, "int64"))
+    )
+    ck.verify(res, 2)
+
+
+def test_simplify_normalize_min_value_expr():
+    ck = CanonicalChecker()
+    x = tvm.tirx.Var("x", "int32")
+
+    ck.verify(tvm.tirx.min_value("int32") - x == 0, x == tvm.tirx.min_value("int32"))
+    ck.verify(tvm.tirx.min_value("int32") + x == 0, tirx.const(False))
+    ck.verify(0 == tvm.tirx.min_value("int32") - x, x == tvm.tirx.min_value("int32"))
+    ck.verify(0 == tvm.tirx.min_value("int32") + x, tirx.const(False))
+    ck.verify(-x + tvm.tirx.min_value("int32") == 0, x == tvm.tirx.min_value("int32"))
+    ck.verify(x + tvm.tirx.min_value("int32") == 0, tirx.const(False))
+    ck.verify(0 == -x + tvm.tirx.min_value("int32"), x == tvm.tirx.min_value("int32"))
+    ck.verify(0 == x + tvm.tirx.min_value("int32"), tirx.const(False))
+
+
+def test_proddiv_simplify():
+    ck = CanonicalChecker()
+    flm = tvm.tirx.floormod
+    fld = tvm.tirx.floordiv
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("y", "int32")
+
+    ck.verify(flm(x * 32 * x, x), 0)
+    ck.verify(flm(z * x * 32 * x * y, x * z), 0)
+    ck.verify(flm(z * x * 32 * x * y, x * z * y * 8 * x), 0)
+    ck.verify(flm(z * x * 32 * (x * y), 6 * x * z), flm(x * y * 16, 3) * (x * z * 2))
+    ck.verify(flm(x * 32 * x, x * z), flm(x * 32, z) * x)
+
+    ck.verify(tmod(x * 32 * x, x), 0)
+    ck.verify(tmod(z * x * 32 * x * y, x * z), 0)
+    ck.verify(tmod(z * x * 32 * (x * y), 6 * x * z), tmod(x * y * 16, 3) * (x * z * 2))
+    ck.verify(tmod(x * 32 * x, x * z), tmod(x * 32, z) * x)
+
+    ck.verify(fld(x * 2 * x * z, 4 * x * x * x), fld(z, x * 2))
+    ck.verify(fld(x * (2 * y) * 3, 3 * y), x * 2)
+    ck.verify(fld(x * (2 * y) * 3, 3 * y * z), fld(x * 2, z))
+
+    ck.verify(tdiv(x * 2 * x * z, 4 * x * x * x), tdiv(z, x * 2))
+    ck.verify(tdiv(x * (2 * y) * 3, 3 * y), x * 2)
+    ck.verify(tdiv(x * (2 * y) * 3, 3 * y * z), tdiv(x * 2, z))
+
+
+def test_floormod_two():
+    ck = CanonicalChecker()
+    flm = tvm.tirx.floormod
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(flm(x * 10 + 1 + y * 2 + 2, 2), 1)
+
+
+def test_simplify_le():
+    ck = CanonicalChecker()
+    # Case 1. Ignore the extra expr if it's small than the division number
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+    ck.analyzer.bind(y, tvm.ir.Range(0, 8))
+    ck.analyzer.bind(z, tvm.ir.Range(0, 2))
+    ck.verify(x * 8 + y < 16, x < 2)
+    ck.verify(x * 8 + z * 4 < 16, x < 2)
+    ck.verify(x * 8 + z * 4 < 16, x < 2)
+
+    # TODO: Not sure why `-2 < x` will be convert to `x > -2`, use a explicit simplify here.
+    ck.verify(x * -8 + y < 16, ck.analyzer.rewrite_simplify(-2 < x))
+    ck.verify(x * -8 + z * 4 < 16, ck.analyzer.rewrite_simplify(-2 < x))
+
+    ck.verify(x * 8 + y + z < 16, x * 8 + y + z < 16)
+
+    n = tvm.tirx.Var("n", "int32")
+    ck.verify(x * 8 + y < n, x * 8 + y < n)
+
+    # Case 2. Simplify the extra expr
+    x1, x2, ty, tx, vec = (
+        tvm.tirx.Var("x1", "int32"),
+        tvm.tirx.Var("x2", "int32"),
+        tvm.tirx.Var("ty", "int32"),
+        tvm.tirx.Var("tx", "int32"),
+        tvm.tirx.Var("vec", "int32"),
+    )
+    ck.analyzer.bind(x1, tvm.ir.Range(0, 2))
+    ck.analyzer.bind(x2, tvm.ir.Range(0, 3))
+    ck.analyzer.bind(ty, tvm.ir.Range(0, 8))
+    ck.analyzer.bind(tx, tvm.ir.Range(0, 32))
+    ck.analyzer.bind(vec, tvm.ir.Range(0, 8))
+    ck.verify(
+        x1 * 5632 + (((x2 * 8 + ty) * 32 + tx) * 8 + vec) % 5632 < 11008,
+        x1 * 22 + (x2 * 8 + ty) % 22 < 43,
+    )
+    ck.verify(tx // 2 % 8 + vec < 8, tx % 16 // 2 + vec < 8)
+
+    # Case 3. No failure
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+    ck.analyzer.bind(y, tvm.ir.Range(0, 1024))
+    ck.verify(x * 1024 + y < z * 7168, x - z * 7 < 0)
+
+
+def test_simplify_le_negative_scale_extra():
+    """Regression: Case 2 of the LT-with-divisible-coeffs rewrite must not
+    fire when the leftover split term has a negative scale.
+
+    The rewrite ``S + xn < 0  ⇔  S/d + xn // d < 0`` is only sound when
+    the leftover ``xn`` has scale ``+1``. With scale ``-1`` the equivalence
+    becomes ``≤`` rather than ``<`` and the rewrite silently strengthens
+    the predicate. The original bug surfaced as ``row > col`` masks of
+    ``.16x*b`` tcgen05 readbacks collapsing to plain ``warp_id > k``
+    comparisons (lower-triangle writes were silently dropped on the
+    boundary warp).
+    """
+    ck = CanonicalChecker()
+    tx = tvm.tirx.Var("tx", "int32")
+    warp = tvm.tirx.Var("warp", "int32")
+    ck.analyzer.bind(tx, tvm.ir.Range(0, 128))
+    ck.analyzer.bind(warp, tvm.ir.Range(0, 4))
+
+    # Same-source joint projection: the comparison genuinely depends on tx
+    # at warp == 0 (e.g. tx == 4 ⇒ 0 < 1 = True; tx == 1 ⇒ 2 < 0 = False),
+    # so the simplifier must keep both sides.  Pre-fix this folded to
+    # ``0 < warp`` and dropped every True case in warp 0.
+    expr = (tx % 4) * 2 < warp * 16 + (tx % 32) // 4
+    ck.verify(expr, expr)
+
+    # The simpler ``scale = -1`` with ``lower_factor = 1`` shape.  Pre-fix
+    # this folded to ``False`` (drops all warp >= 1 cases where the rhs
+    # actually exceeds 8*warp).
+    expr = warp * 8 < (tx % 32)
+    ck.verify(expr, expr)
+
+    # The corresponding ``scale = +1`` Case 2 path (the rewrite this guards)
+    # must still optimize — verifies we did not over-restrict.
+    x1 = tvm.tirx.Var("x1", "int32")
+    y1 = tvm.tirx.Var("y1", "int32")
+    ck.verify(x1 * 64 + (y1 % 64) < 120, x1 * 8 + (y1 % 64) // 8 < 15)
+
+    # The truly-always-true comparison that arises from the same kernel
+    # (``r = 2 / va = 1`` in the tcgen05.ld.16x256b readback) must still
+    # fold to True so the masked store can be elided.
+    expr_true = (tx % 4) * 2 < warp * 16 + (tx % 32) // 4 + 8
+    ck.verify(expr_true, tvm.tirx.const(True, "bool"))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_const_int_bound.py

@@ -0,0 +1,323 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: RUF012
+
+import contextlib
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.arith import ConstIntBound
+
+NEG_INF = ConstIntBound.NEG_INF
+POS_INF = ConstIntBound.POS_INF
+
+
+class TestCase:
+    def __init__(self, expr, expected_bounds, known_bounds=None, constraint=None):
+        self.expr = expr
+        self.expected_bounds = expected_bounds
+        if known_bounds is None:
+            self.known_bounds = {}
+        else:
+            self.known_bounds = known_bounds
+
+        self.constraint = constraint
+
+    @property
+    def __name__(self):
+        return str(self.expr)
+
+
+class BaseCompare:
+    def test_const_bounds(self, test_case):
+        analyzer = tvm.arith.Analyzer()
+
+        for var, bounds in test_case.known_bounds.items():
+            analyzer.update(var, ConstIntBound(*bounds))
+            assert analyzer.const_int_bound_is_bound(var)
+
+        with contextlib.ExitStack() as stack:
+            if test_case.constraint is not None:
+                stack.enter_context(analyzer.constraint_scope(test_case.constraint))
+
+            bounds = analyzer.const_int_bound(test_case.expr)
+
+        if test_case.expected_bounds[0] is None:
+            assert bounds.max_value == test_case.expected_bounds[1]
+        elif test_case.expected_bounds[1] is None:
+            assert bounds.min_value == test_case.expected_bounds[0]
+        else:
+            assert (bounds.min_value, bounds.max_value) == test_case.expected_bounds
+
+
+class TestDataType(BaseCompare):
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.Var("x", "int64"), (NEG_INF, POS_INF)),
+        TestCase(tvm.tirx.Var("x", "int8"), (-128, 127)),
+        TestCase(tvm.tirx.Var("x", "uint8"), (0, 255)),
+        TestCase(tvm.tirx.Var("x", "int32"), (-(2**31), 2**31 - 1)),
+    )
+
+
+def test_plain_var_non_negative_bound_requires_context():
+    var = tvm.tirx.Var("x", "int64")
+    analyzer = tvm.arith.Analyzer()
+
+    assert analyzer.const_int_bound(var).min_value == NEG_INF
+    with analyzer.constraint_scope(var >= 0):
+        assert analyzer.const_int_bound(var).min_value == 0
+    assert analyzer.const_int_bound(var).min_value == NEG_INF
+
+
+class TestCastBound(BaseCompare):
+    x = tvm.tirx.Var("x", "int8")
+    tmod = tvm.tirx.truncmod
+
+    test_case = tvm.testing.parameter(
+        TestCase(tmod(x, 3).astype("uint32"), (0, 2)),
+        TestCase(tmod(x, 3).astype("float32").astype("int32"), (-2, 2)),
+    )
+
+
+class TestAddSubBound(BaseCompare):
+    x = tvm.tirx.Var("x", "int64")
+    y = tvm.tirx.Var("y", "int64")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x + y, (NEG_INF, POS_INF)),
+        TestCase(x + y, (1, 14), known_bounds={x: (0, 4), y: (1, 10)}),
+        TestCase(x - y, (-10, 3), known_bounds={x: (0, 4), y: (1, 10)}),
+        TestCase(x - y, (-10, POS_INF), known_bounds={x: (0, POS_INF), y: (1, 10)}),
+        TestCase(1 - x, (NEG_INF, 1), known_bounds={x: (0, POS_INF), y: (1, 10)}),
+    )
+
+
+@pytest.mark.xfail(reason="Not currently supported")
+class TestBoundsUsingReciprocals(BaseCompare):
+    """Special handling for differences of reciprocals
+
+    These terms can appear when comparing the number of operations for
+    different orderings of matrix multiplications, with A, B, and C
+    known to be positive values.
+
+    In these cases, comparing `(A+B)*C < A*B` is equivalent to
+    `1/A + 1/B < 1/C`.  Working in terms of the reciprocals
+    allows the ConstIntBound analyzer to provide a tighter
+    bound for these differences than would otherwise be
+    available.
+
+    For `(A+B)*C - A*B`, the normal bottom-up integer bounds are unable to
+    provide the bounds required to provide these inequalities, because they
+    treat the terms as uncorrelated.  That is, they assume that `(A+B)*C` may
+    achieve its minimum while `A*B` simultaneously achieves its maximum.
+    """
+
+    A, B, C = [tvm.tirx.Var(letter, "int64") for letter in "ABC"]
+
+    symmetric_bounds = {A: (1, 4095), B: (1, 4095), C: (2048, 2048)}
+    asymmetric_bounds = {A: (1, 1024), B: (1, POS_INF), C: (2048, 2048)}
+
+    test_case = tvm.testing.parameter(
+        TestCase((A + B) * C - A * B, (2048, None), known_bounds=symmetric_bounds),
+        TestCase((A + B) * C - B * A, (2048, None), known_bounds=symmetric_bounds),
+        TestCase(A * B - (A + B) * C, (None, -2048), known_bounds=symmetric_bounds),
+        TestCase(B * A - (A + B) * C, (None, -2048), known_bounds=symmetric_bounds),
+        TestCase((A + B) * C - A * B, (2048, None), known_bounds=asymmetric_bounds),
+        TestCase((A + B) * C - B * A, (2048, None), known_bounds=asymmetric_bounds),
+        TestCase(A * B - (A + B) * C, (None, -2048), known_bounds=asymmetric_bounds),
+        TestCase(B * A - (A + B) * C, (None, -2048), known_bounds=asymmetric_bounds),
+    )
+
+
+class TestMulBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x * y + 20, (0, 60), {x: (-2, 4), y: (4, 10)}),
+        TestCase(x * y, (-32, 24), {x: (-3, 4), y: (-8, 2)}),
+        TestCase(x * y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-8, 2)}),
+    )
+
+
+class TestTruncDivBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    expr = tvm.tirx.truncdiv(x, y)
+
+    test_case = tvm.testing.parameter(
+        TestCase(expr, (-2, None), {x: (-9, 4), y: (4, 10)}),
+        TestCase(expr, (-4, 9), {x: (-9, 4), y: (-2, 0)}),
+        TestCase(expr, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}),
+        TestCase(expr, (-9, 9), {x: (-9, 4), y: (-4, 12)}),
+    )
+
+
+class TestTruncModBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    expr = tvm.tirx.truncmod(x, y)
+
+    test_case = tvm.testing.parameter(
+        TestCase(expr, (-9, 4), {x: (-9, 4), y: (4, 10)}),
+        TestCase(expr, (-9, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}),
+        TestCase(expr, (0, 9), {x: (1, POS_INF), y: (4, 10)}),
+    )
+
+
+class TestFloorDivBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ux = tvm.tirx.Var("x", "uint32")
+    uy = tvm.tirx.Var("y", "uint32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x // y, (-9 // 4, None), {x: (-9, 4), y: (4, 10)}),
+        TestCase(x // y, (-4, 9), {x: (-9, 4), y: (-2, 0)}),
+        TestCase(x // y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}),
+        TestCase(x // y, (-9, 9), {x: (-9, 4), y: (-4, 12)}),
+        TestCase(ux // uy, (0, 4), {ux: (1, 4), uy: (0, 12)}),
+    )
+
+
+class TestFloorModBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x % y, (0, 9), {x: (-9, 4), y: (4, 10)}),
+        TestCase(x % y, (0, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}),
+        TestCase(x % y, (0, 9), {x: (1, POS_INF), y: (4, 10)}),
+    )
+
+
+class TestMinMaxBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.min(x, y), (-9, 10), {x: (-9, 11), y: (4, 10)}),
+        TestCase(tvm.tirx.min(x, y), (NEG_INF, 10), {x: (NEG_INF, POS_INF), y: (4, 10)}),
+        TestCase(tvm.tirx.max(x, y), (4, POS_INF), {x: (NEG_INF, POS_INF), y: (4, 10)}),
+        TestCase(tvm.tirx.max(x, y), (4, POS_INF), {x: (1, POS_INF), y: (4, 10)}),
+    )
+
+
+class TestSelectBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(
+            tvm.tirx.Select(x > 1, (y < 0).astype("int32"), y + 1),
+            (0, 11),
+            {x: (-9, 11), y: (4, 10)},
+        ),
+    )
+
+
+class TestShiftAndBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(x >> y, (-3, 2), {x: (-9, 11), y: (2, 10)}),
+        TestCase(x & y, (0, 10), {x: (-9, 11), y: (2, 10)}),
+        TestCase(x & y, (0, 10), {x: (10, 11), y: (2, 10)}),
+    )
+
+
+class TestMixIndexBound(BaseCompare):
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    tdiv = tvm.tirx.truncdiv
+    tmod = tvm.tirx.truncmod
+
+    test_case = tvm.testing.parameter(
+        TestCase(tmod(x, 8) + tdiv(x, 8) * 8, (0, 24 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}),
+        TestCase(y + x * 3, (0, 24 * 3 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}),
+        TestCase(
+            tmod(x, 7) + tdiv(x, 7) * 7, (0, (23 // 7) * 7 + 6), {x: (0, 24 - 1), y: (0, 3 - 1)}
+        ),
+    )
+
+
+class TestLetBound(BaseCompare):
+    x = tvm.tirx.Var("x", "int32")
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.Let(x, 1, x + 1), (2, 2)),
+    )
+
+
+class TestFloorModNegativeDivisor(BaseCompare):
+    flm, fld = tvm.tirx.floormod, tvm.tirx.floordiv
+    a, b = tvm.tirx.Var("a", "int32"), tvm.tirx.Var("b", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(a % b, (-4, 6), {a: (0, 6), b: (-5, 7)}),
+    )
+
+
+class TestDivModAssumeNoZeroDivisor(BaseCompare):
+    """Divmod non negative expression makes assumption that divide by
+    zero won't occur this assumption is important to get best result
+    from symbolic shape programs
+    """
+
+    a, b = tvm.tirx.Var("a", "int32"), tvm.tirx.Var("b", "int32")
+
+    test_case = tvm.testing.parameter(
+        TestCase(a // b, (0, 6), {a: (0, 6), b: (0, POS_INF)}),
+        TestCase(a % b, (0, 6), {a: (0, 6), b: (0, POS_INF)}),
+    )
+
+
+class TestMultipleCondition(BaseCompare):
+    a = tvm.tirx.Var("a", "int32")
+    test_case = tvm.testing.parameter(
+        TestCase(
+            a % 58 - 1,
+            (0, None),
+            known_bounds={a: (0, 128)},
+            constraint=tvm.tirx.all(1 <= a % 58, a % 58 < 57),
+        ),
+    )
+
+
+class TestBroadcastBound(BaseCompare):
+    a = tvm.tirx.Var("a", "int32")
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.Broadcast(a, 4), (0, 128), {a: (0, 128)}),
+    )
+
+
+class TestRampBound(BaseCompare):
+    a = tvm.tirx.Var("a", "int32")
+    test_case = tvm.testing.parameter(
+        TestCase(tvm.tirx.Ramp(a, 2, 4) + 2, (2, 128 + 2 * 3 + 2), {a: (0, 128)}),
+    )
+
+
+class TestModularSetBound(BaseCompare):
+    analyzer = tvm.arith.Analyzer()
+    tx = tvm.tirx.Var("tx", "int32")
+    bx = tvm.tirx.Var("bx", "int32")
+
+    expr = (bx * 2048 + tx * 16) % 7168
+
+    test_case = tvm.testing.parameter(
+        TestCase(expr, (0, 7152), {bx: (0, 3584), tx: (0, 128)}),
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_deduce_bound.py

@@ -0,0 +1,272 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401, F841
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.tirx.buffer import decl_buffer
+
+
+def test_deduce():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    c = tvm.tirx.Var("c", "int32")
+    d = tvm.tirx.Var("d", "int32")
+
+    b_s = tvm.arith.IntervalSet(2, 3)
+    c_s = tvm.arith.IntervalSet(10, 15)
+    d_s = tvm.arith.IntervalSet(-3, -1)
+    zero = tvm.tirx.const(0, "int32")
+
+    fdiv = tvm.tirx.floordiv
+
+    e0 = (-b) * a + c - d
+    res0 = tvm.arith.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {})
+    ans0 = fdiv(d - c, b * -1)
+    tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+
+    # expression containing variable a is on rhs
+    res0 = tvm.arith.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {})
+    tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+
+    e0 = d * a + c - d
+    res0 = tvm.arith.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {})
+    ans0 = fdiv(d - c, d)
+    tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+
+    # expression containing variable a is on rhs
+    res0 = tvm.arith.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {})
+    tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+
+    e1 = a * 4 + b < c
+    res1 = tvm.arith.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {})
+    ans1 = fdiv(c - 1 - b, 4)
+    tvm.testing.assert_prim_expr_equal(res1.max_value, ans1)
+
+    # expression containing variable a is on rhs
+    e1 = c > a * 4 + b
+    res1 = tvm.arith.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {})
+    tvm.testing.assert_prim_expr_equal(res1.max_value, ans1)
+
+    e2 = tvm.tirx.max(5, a * 4) < 0
+    res2 = tvm.arith.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {})
+    assert str(res2.max_value) == "neg_inf"
+    assert str(res2.min_value) == "pos_inf"
+
+    # expression containing variable a is on rhs
+    e2 = zero < tvm.tirx.max(5, a * 4)
+    res2 = tvm.arith.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {})
+    assert str(res2.max_value) == "neg_inf"
+    assert str(res2.min_value) == "pos_inf"
+
+    e3 = (-b) + a * c - d
+    res3 = tvm.arith.deduce_bound(a, e3 >= 0, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s})
+    ans3 = fdiv(2, c) + 1
+    tvm.testing.assert_prim_expr_equal(res3.min_value, ans3)
+
+    res3 = tvm.arith.deduce_bound(a, zero <= e3, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s})
+    tvm.testing.assert_prim_expr_equal(res3.min_value, ans3)
+
+    # tests for `EQ` op
+    res4 = tvm.arith.deduce_bound(a, a == b, {}, {})
+    tvm.testing.assert_prim_expr_equal(res4.max_value, b)
+    tvm.testing.assert_prim_expr_equal(res4.min_value, b)
+
+    # Unsatisfiable `EQ`, variable as one of the Operand
+    res5 = tvm.arith.deduce_bound(a, (a == b), {b: b_s}, {b: b_s})
+    assert str(res5.max_value) == "neg_inf"
+    assert str(res5.min_value) == "pos_inf"
+
+    # variable `a` on the RHS side
+    res6 = tvm.arith.deduce_bound(a, 10 == a, {}, {})
+    tvm.testing.assert_prim_expr_equal(res6.max_value, 10)
+    tvm.testing.assert_prim_expr_equal(res6.min_value, 10)
+
+    # Add, Sub in `EQ`
+    e4 = (a - c) == (b + d)
+    ans4 = b + d + c
+    res7 = tvm.arith.deduce_bound(a, e4, {b: b_s, c: c_s, d: d_s}, {})
+    tvm.testing.assert_prim_expr_equal(res7.max_value, ans4)
+    tvm.testing.assert_prim_expr_equal(res7.min_value, ans4)
+
+    # Satisfiable Mul in `EQ` with negative sign
+    res8 = tvm.arith.deduce_bound(a, (5 * a == -10), {}, {})
+    tvm.testing.assert_prim_expr_equal(res8.max_value, -2)
+    tvm.testing.assert_prim_expr_equal(res8.min_value, -2)
+
+    # Unsatisfiable Mul in `EQ`
+    e5 = 4 * a == b
+    res9 = tvm.arith.deduce_bound(a, e5, {b: b_s}, {})
+    assert str(res9.max_value) == "neg_inf"
+    assert str(res9.min_value) == "pos_inf"
+
+    res10 = tvm.arith.deduce_bound(a, (b * a == b), {b: b_s}, {})
+    # simplifier is now able to prove symbolic relation (b * a % b == 0)
+    tvm.testing.assert_prim_expr_equal(res10.max_value, 1)
+    tvm.testing.assert_prim_expr_equal(res10.min_value, 1)
+
+
+def test_check():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    c = tvm.tirx.Var("c", "int32")
+    d = tvm.tirx.Var("d", "int32")
+
+    b_s = tvm.arith.IntervalSet(2, 3)
+    c_s = tvm.arith.IntervalSet(5, 7)
+    d_s = tvm.arith.IntervalSet(-3, -1)
+
+    # no compare operator
+    res1 = tvm.arith.deduce_bound(a, a + b, {b: b_s}, {})
+    assert res1.is_nothing()
+
+    # multiple compare operators
+    res2 = tvm.arith.deduce_bound(a, (a + b > 3).astype(c.ty) > c, {b: b_s, c: c_s}, {})
+    assert res2.is_nothing()
+
+    # multiple target variable
+    res2 = tvm.arith.deduce_bound(a, a * 2 - a > b, {b: b_s}, {})
+    assert res2.is_nothing()
+
+
+def test_deduce_basic():
+    def test_basic(a1, a2, coff):
+        a = tvm.tirx.Var("a", "int32")
+        b = tvm.tirx.Var("b", "int32")
+        b_s = tvm.arith.IntervalSet(a1, a2)
+        e0 = b + a * coff + 3
+
+        res1 = tvm.arith.deduce_bound(a, e0 < 17, {b: b_s}, {b: b_s})
+        [x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) < 17, True)
+
+        # expression containing variable a is on rhs
+        res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(17, "int32") < e0, {b: b_s}, {b: b_s})
+        [x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) > 17, True)
+
+        # expression containing variable a is on rhs
+        res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(17, "int32") >= e0, {b: b_s}, {b: b_s})
+        [x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
+
+        tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) <= 17, True)
+
+        res1 = tvm.arith.deduce_bound(a, e0 >= 17, {b: b_s}, {b: b_s})
+        [x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) >= 17, True)
+
+    test_basic(0, 4, 4)
+    test_basic(1, 5, 4)
+    test_basic(2, 6, 4)
+    test_basic(0, 4, -4)
+    test_basic(1, 5, -4)
+    test_basic(2, 6, -4)
+
+
+def test_deduce_complex():
+    def test_complex(a1, a2, coff):
+        a = tvm.tirx.Var("a", "int32")
+        b = tvm.tirx.Var("b", "int32")
+        b_s = tvm.arith.IntervalSet(a1, a2)
+        e0 = (b * 3 + a * coff) * 4
+
+        res1 = tvm.arith.deduce_bound(a, e0 < 63, {b: b_s}, {b: b_s})
+        [t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) < 63, True)
+
+        # expression containing variable a is on rhs
+        res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(63, "int32") >= e0, {b: b_s}, {b: b_s})
+        [t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) <= 63, True)
+
+        res1 = tvm.arith.deduce_bound(a, e0 > 63, {b: b_s}, {b: b_s})
+        [t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) > 63, True)
+
+        # expression containing variable a is on rhs
+        res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(63, "int32") <= e0, {b: b_s}, {b: b_s})
+        [t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
+        tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) >= 63, True)
+
+    test_complex(0, 4, 4)
+    test_complex(0, 4, -4)
+    test_complex(2, 6, 4)
+    test_complex(0, 4, -4)
+    test_complex(1, 5, -4)
+    test_complex(2, 6, -4)
+
+
+def test_deduce_non_support():
+    a = tvm.tirx.Var("a", "int32")
+
+    def test_non_support(lhs):
+        res = tvm.arith.deduce_bound(a, lhs < 10, {}, {})
+        assert res.is_nothing()
+
+    test_non_support(tvm.tirx.floormod(a, 16))
+    test_non_support(tvm.tirx.Min(a, 16))
+    test_non_support(tvm.tirx.Max(a, 16))
+    test_non_support(tvm.tirx.LE(a, 16))
+    test_non_support(tvm.tirx.LT(a, 16))
+    test_non_support(tvm.tirx.GE(a, 16))
+    test_non_support(tvm.tirx.GT(a, 16))
+    test_non_support(tvm.tirx.EQ(a, 16))
+    test_non_support(tvm.tirx.NE(a, 16))
+    test_non_support(tvm.tirx.log(a))
+    test_non_support(tvm.tirx.BufferLoad(decl_buffer([16], "int32"), [a]))
+
+
+def test_deduce_floordiv():
+    def do_test(gen_expr, dom_map, expect_min, expect_max):
+        a = tvm.tirx.Var("a", "int32")
+        expr = gen_expr(a)
+        res = tvm.arith.deduce_bound(a, expr, dom_map, dom_map)
+        if isinstance(expect_min, str):
+            assert str(res.min_value) == expect_min
+        else:
+            tvm.testing.assert_prim_expr_equal(res.min_value, expect_min)
+        if isinstance(expect_max, str):
+            assert str(res.max_value) == expect_max
+        else:
+            tvm.testing.assert_prim_expr_equal(res.max_value, expect_max)
+
+    # test basic cases
+    do_test(lambda a: a // 8 > 3, {}, 32, "pos_inf")
+    do_test(lambda a: a // 8 >= 3, {}, 24, "pos_inf")
+    do_test(lambda a: a // 8 < 3, {}, "neg_inf", 23)
+    do_test(lambda a: a // 8 <= 3, {}, "neg_inf", 31)
+    do_test(lambda a: a // 8 == 3, {}, "pos_inf", "neg_inf")
+    do_test(lambda a: a // 8 > -3, {}, -16, "pos_inf")
+    do_test(lambda a: a // 8 >= -3, {}, -24, "pos_inf")
+    do_test(lambda a: a // -8 > 3, {}, "neg_inf", -32)
+    do_test(lambda a: a // -8 >= 3, {}, "neg_inf", -24)
+    do_test(lambda a: a // -8 < 3, {}, -23, "pos_inf")
+    do_test(lambda a: a // -8 <= 3, {}, -31, "pos_inf")
+    do_test(lambda a: 8 // a >= 2, {}, "pos_inf", "neg_inf")
+
+    # test nested cases
+    b = tvm.tirx.Var("b", "int32")
+    bs = {b: tvm.arith.IntervalSet(2, 6)}
+    do_test(lambda a: b * 3 + a // 8 < 63, bs, "neg_inf", 359)
+    do_test(lambda a: b * 3 + a // 8 <= 63, bs, "neg_inf", 367)
+    do_test(lambda a: b * 3 + a // 8 > 63, bs, 464, "pos_inf")
+    do_test(lambda a: b * 3 + a // 8 >= 63, bs, 456, "pos_inf")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_detect_clip_bound.py

@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+import tvm.testing
+
+
+def test_basic():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    c = tvm.tirx.Var("c", "int32")
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a])
+    tvm.testing.assert_prim_expr_equal(m[1], b * 6 - 1)
+    assert m[0].value == 2
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a, b])
+    assert len(m) == 0
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(a + 10 * c <= 20, b - 1 > 0), [a, b])
+    tvm.testing.assert_prim_expr_equal(m[1], 20 - 10 * c)
+    tvm.testing.assert_prim_expr_equal(m[2], 2)
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(tvm.tirx.Not(a * 1 > b * 6), a - 1 > 0), [a])
+    tvm.testing.assert_prim_expr_equal(m[1], b * 6)
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(tvm.tirx.Min(a, b) > 3, a - 10 < 0), [a, b])
+    tvm.testing.assert_prim_expr_equal(m[0], 4)
+    tvm.testing.assert_prim_expr_equal(m[1], 9)
+    tvm.testing.assert_prim_expr_equal(m[2], 4)
+
+
+def test_trivial_eq():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    m = tvm.arith.detect_clip_bound(b == 3, [a, b])
+    tvm.testing.assert_prim_expr_equal(m[2], 3)
+    tvm.testing.assert_prim_expr_equal(m[3], 3)
+    m = tvm.arith.detect_clip_bound(tvm.tirx.all(a == 4, b == 3), [a, b])
+    tvm.testing.assert_prim_expr_equal(m[0], 4)
+    tvm.testing.assert_prim_expr_equal(m[1], 4)
+    tvm.testing.assert_prim_expr_equal(m[2], 3)
+    tvm.testing.assert_prim_expr_equal(m[3], 3)
+
+
+if __name__ == "__main__":
+    test_basic()

tests/python/arith/test_arith_detect_linear_equation.py

@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+import tvm.testing
+
+
+def test_basic():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    m = tvm.arith.detect_linear_equation(a * 4 + b * 6 + 7, [a])
+    assert m[0].value == 4
+    tvm.testing.assert_prim_expr_equal(m[1], b * 6 + 7)
+
+    m = tvm.arith.detect_linear_equation(a * 4 * (a + 1) + b * 6 + 7, [a])
+    assert len(m) == 0
+
+    m = tvm.arith.detect_linear_equation(a * 4 + (a + 1) + b * 6 + 7, [a])
+    assert m[0].value == 5
+    tvm.testing.assert_prim_expr_equal(m[1], b * 6 + 7 + 1)
+
+    m = tvm.arith.detect_linear_equation(a * b + 7, [a])
+    assert m[0] == b
+
+    m = tvm.arith.detect_linear_equation(b * 7, [a])
+    assert m[0].value == 0
+
+    m = tvm.arith.detect_linear_equation(b * 7, [])
+    assert len(m) == 1
+    tvm.testing.assert_prim_expr_equal(m[0], b * 7)
+
+    c = tvm.tirx.Var("c", "uint32")
+    m = tvm.arith.detect_linear_equation(128 - c, [c])
+    assert m[0].value == -1
+
+
+def test_multivariate():
+    v = [tvm.tirx.Var(f"v{i}", "int32") for i in range(4)]
+    b = tvm.tirx.Var("b", "int32")
+    m = tvm.arith.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8, v)
+
+    tvm.testing.assert_prim_expr_equal(m[0], b + 5)
+
+    assert m[1].value == 8
+
+    m = tvm.arith.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8 * v[2], v)
+    assert len(m) == 0
+
+    m = tvm.arith.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8 * v[1] + v[3], v)
+    assert len(m) == 0
+
+    m = tvm.arith.detect_linear_equation(((v[0] * b + v[1]) * 8 + v[2] + 1) * 2, v)
+    assert m[1].value == 16
+    assert m[2].value == 2
+    assert m[len(m) - 1].value == 2
+
+    m = tvm.arith.detect_linear_equation((v[0] - v[1]), [v[2]])
+    assert m[0].value == 0
+
+    tvm.testing.assert_prim_expr_equal(m[1], v[0] - v[1])
+
+    m = tvm.arith.detect_linear_equation((v[0] - v[1]), [])
+    assert len(m) == 1
+    tvm.testing.assert_prim_expr_equal(m[0], v[0] - v[1])
+
+
+if __name__ == "__main__":
+    test_basic()
+    test_multivariate()

tests/python/arith/test_arith_domain_touched.py

@@ -0,0 +1,95 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pytest
+import tvm_ffi
+
+import tvm
+from tvm.script import tirx as T
+
+
+@T.prim_func(s_tir=True)
+def scalar_func(a: T.handle, b: T.handle):
+    m = T.int32()
+    n = T.meta_var(100)
+    A = T.match_buffer(a, (n, m))
+    B = T.match_buffer(b, (n, m))
+
+    for i, j in T.grid(n, m):
+        A[i, j] = B[i - 1, j + 1] + A[i - 1, j - 1]
+
+
+def test_domain_touched():
+    func = scalar_func
+    a, b = [func.buffer_map[var] for var in func.params]
+    ir = func.body
+
+    a_domain_r = tvm.arith._ffi_api.DomainTouched(ir, a, True, False)
+
+    assert a_domain_r[0].min.value == -1
+    assert a_domain_r[0].extent.value == 100
+    assert a_domain_r[1].min.value == -1
+    assert a_domain_r[1].extent.name == "m"
+
+    a_domain_w = tvm.arith._ffi_api.DomainTouched(ir, a, False, True)
+    assert a_domain_w[0].min.value == 0
+    assert a_domain_w[0].extent.value == 100
+    assert a_domain_w[1].min.value == 0
+    assert a_domain_w[1].extent.name == "m"
+
+    a_domain_rw = tvm.arith._ffi_api.DomainTouched(ir, a, True, True)
+    assert a_domain_rw[0].min.value == -1
+    assert a_domain_rw[0].extent.value == 101
+    assert a_domain_rw[1].min.value == -1
+    assert isinstance(a_domain_rw[1].extent, tvm.tirx.Add)
+    assert a_domain_rw[1].extent.a.name == "m"
+    assert a_domain_rw[1].extent.b.value == 1
+
+    b_domain_r = tvm.arith._ffi_api.DomainTouched(ir, b, True, False)
+    assert b_domain_r
+    assert b_domain_r[0].min.value == -1
+    assert b_domain_r[0].extent.value == 100
+    assert b_domain_r[1].min.value == 1
+    assert b_domain_r[1].extent.name == "m"
+
+    b_domain_w = tvm.arith._ffi_api.DomainTouched(ir, b, False, True)
+    assert isinstance(b_domain_w, tvm_ffi.Array)
+    assert len(b_domain_w) == 0
+
+
+def test_domain_touched_vector():
+    pytest.skip("BufferRegion arithmetic in expressions not supported")
+    m = tvm.runtime.convert(128)
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle, n: T.int32):
+        A = T.match_buffer(a, (n * m,))
+        B = T.match_buffer(b, (n * m,))
+
+        for i in T.serial(n):
+            A[i * m : (i + 1) * m : 1] = A[i * m : (i + 1) * m : 1] + B[i * m : (i + 1) * m : 1]
+
+    a, b = [func.buffer_map[var] for var in func.params[:2]]
+
+    assert tvm.arith._ffi_api.DomainTouched(func.body, a, True, False)[0].extent.value == 128
+    assert tvm.arith._ffi_api.DomainTouched(func.body, a, True, False)[0].extent.value == 128
+    assert tvm.arith._ffi_api.DomainTouched(func.body, a, True, True)[0].extent.value == 128
+    assert tvm.arith._ffi_api.DomainTouched(func.body, b, True, False)[0].extent.value == 128
+    assert tvm.arith._ffi_api.DomainTouched(func.body, b, True, False)[0].extent.value == 128
+
+
+if __name__ == "__main__":
+    test_domain_touched()

tests/python/arith/test_arith_intset.py

@@ -0,0 +1,456 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.arith.analyzer import Analyzer
+
+
+class IntSetChecker:
+    def __init__(self):
+        self.analyzer = tvm.arith.Analyzer()
+
+    def verify(self, data, dmap, expected):
+        res = self.analyzer.int_set(data, dmap)
+
+        def err_msg():
+            return f"\ndata={data}\ndmap={dmap}\nres={res}\nexpected={expected}"
+
+        assert self.analyzer.can_prove_equal(res.min_value, expected[0]), err_msg()
+        assert self.analyzer.can_prove_equal(res.max_value, expected[1]), err_msg()
+
+
+def test_basic():
+    s = tvm.arith.IntervalSet(2, 3)
+    assert s.min_value.value == 2
+    assert s.max_value.value == 3
+
+    s = tvm.arith.IntSet.single_point(2)
+    assert s.min_value.value == 2
+    assert s.max_value.value == 2
+
+
+def test_vector():
+    base = 10
+    stride = 3
+    lanes = 2
+    s = tvm.arith.IntSet.vector(tvm.tirx.Ramp(base, stride, lanes))
+    assert s.min_value.value == base
+    assert s.max_value.value == base + stride * (lanes - 1)
+
+
+def test_scalable_vector():
+    base = 5
+    s = tvm.arith.IntSet.vector(tvm.tirx.Ramp(base, 2, tvm.tirx.vscale() * 4))
+
+    assert s.min_value.value == base
+    assert s.max_value.same_as(tvm.arith.int_set.pos_inf())
+
+
+def test_add_sub():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(x + y, {x: tvm.arith.IntervalSet(0, 10)}, (y, 10 + y))
+    ck.verify(x + y, {x: tvm.arith.IntervalSet(0, 10), y: tvm.arith.IntervalSet(1, 11)}, (1, 21))
+    ck.verify(x - y, {x: tvm.arith.IntervalSet(0, 10), y: tvm.arith.IntervalSet(1, 11)}, (-11, 9))
+
+
+def test_mul_div():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    tdiv = tvm.tirx.truncdiv
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(1, 100), override=True)
+    ck.verify(x * y, {x: tvm.arith.IntervalSet(0, 10)}, (0, 10 * y))
+    ck.verify(x * 2, {x: tvm.arith.IntervalSet(1, 10)}, (2, 20))
+    ck.verify(x * -2, {x: tvm.arith.IntervalSet(1, 10)}, (-20, -2))
+
+    ck.verify(tdiv(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, tdiv(10, y)))
+    ck.verify(tdiv(x, 2), {x: tvm.arith.IntervalSet(1, 10)}, (0, 5))
+
+    fld = tvm.tirx.floordiv
+    ck.verify(fld(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, fld(10, y)))
+    ck.verify(fld(x, 2), {x: tvm.arith.IntervalSet(-1, 10)}, (-1, 5))
+
+
+def test_mod():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    tmod = tvm.tirx.truncmod
+    ck.analyzer.update(y, tvm.arith.ConstIntBound(1, 100), override=True)
+    ck.verify(tmod(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, y - 1))
+    ck.verify(tmod(x, 10), {x: tvm.arith.IntervalSet(1, 10)}, (0, 9))
+
+    flm = tvm.tirx.floormod
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(-10, 10)}, (0, 9))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 5)}, (3, 5))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(13, 15)}, (3, 5))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 15)}, (0, 9))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 11)}, (0, 9))
+    ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(1, 21)}, (0, 9))
+
+    fld = tvm.tirx.floordiv
+    z = tvm.tirx.Var("z", "int32")
+    ck.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 3))
+    ck.verify(
+        flm(y, 8),
+        {y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)},
+        (
+            z * 8 + x * 4 - 8 * fld(z * 8 + x * 4, 8),
+            z * 8 + x * 4 + 3 - 8 * fld(z * 8 + x * 4, 8),
+        ),
+    )
+    ck1 = IntSetChecker()
+    ck1.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 2))
+    ck1.verify(
+        flm(y, 8), {y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)}, (x * 4, x * 4 + 3)
+    )
+
+
+def test_max_min():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(tvm.tirx.max(x, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (1, 11))
+    ck.verify(tvm.tirx.min(x - 1, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (-1, 9))
+    ck.verify(tvm.tirx.min(x, y), {}, (tvm.tirx.min(x, y), tvm.tirx.min(x, y)))
+    ck.verify(tvm.tirx.max(x, y), {}, (tvm.tirx.max(x, y), tvm.tirx.max(x, y)))
+
+
+def test_select():
+    ck = IntSetChecker()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ck.verify(tvm.tirx.Select(x > 0, x - 1, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (-1, 11))
+
+
+def check_region_bound(expect_region, var_dom, mode, predicate=None):
+    """Helper to check region bound estimation.
+
+    Parameters
+    ----------
+    expect_region: dict
+        The keys are of form (begin, end) or Expr as a single point. The values are
+        expected estimated region or region dict on different bindings.
+
+    var_dom: dict
+        Map var to iteration domain range.
+
+    mode: str
+        Specify "lowerbound", "upperbound" or else use strict bound estimation.
+
+    predicate: Expr
+        Extra predicate, defaults to True.
+    """
+    if predicate is None:
+        predicate = tvm.tirx.IntImm("bool", 1)
+    region = []
+    expect = []
+    for k, v in expect_region.items():
+        if not isinstance(k, tuple | list):
+            k = (k, k + 1)
+        region.append(tvm.ir.Range.from_min_extent(k[0], Analyzer().simplify(k[1] - k[0])))
+        expect.append(v)
+    if mode == "lowerbound":
+        result = tvm.arith.estimate_region_lower_bound(
+            region=region, var_dom=var_dom, predicate=predicate
+        )
+    elif mode == "upperbound":
+        result = tvm.arith.estimate_region_upper_bound(
+            region=region, var_dom=var_dom, predicate=predicate
+        )
+    else:
+        result = tvm.arith.estimate_region_strict_bound(
+            region=region, var_dom=var_dom, predicate=predicate
+        )
+    if result is None:
+        assert all([_ is None for _ in expect])
+        return
+    assert len(result) == len(expect)
+    for intset, expect_desc in zip(result, expect):
+        if isinstance(expect_desc, dict):
+            # check range on different free var bindings
+            for binding in expect_desc:
+                analyzer = Analyzer()
+                for k, v in binding:
+                    analyzer.bind(k, v)
+                expect_begin, expect_end = expect_desc[binding]
+                result_begin = analyzer.simplify(intset.min_value, 3)
+                result_end = analyzer.simplify(intset.max_value + 1, 3)
+                assert analyzer.can_prove_equal(result_begin - expect_begin, 0), (
+                    f"{result_begin} vs {expect_begin}"
+                )
+                assert analyzer.can_prove_equal(result_end - expect_end, 0), (
+                    f"{result_end} vs {expect_end}"
+                )
+        else:
+            # check range
+            expect_begin, expect_end = expect_desc
+            analyzer = Analyzer()
+            assert analyzer.can_prove_equal(intset.min_value - expect_begin, 0), (
+                f"{intset.min_value} vs {expect_begin}"
+            )
+            assert analyzer.can_prove_equal(intset.max_value - expect_end + 1, 0), (
+                f"{intset.max_value} vs {expect_end - 1}"
+            )
+
+
+def test_region_bound_not_independent():
+    # (i, i+2) and (i+2, i+4) are dependent, this the lowerbound is not available
+    i = tvm.tirx.Var("i", "int32")
+    var_dom = {
+        i: tvm.ir.Range(begin=0, end=64),
+    }
+    check_region_bound({(i, i + 2): None, (i + 2, i + 4): None}, var_dom, mode="lowerbound")
+    check_region_bound({(i, i + 2): (0, 65), (i + 2, i + 4): (2, 67)}, var_dom, mode="upperbound")
+
+    # when only a subset of access indices are affine
+    i, j, k = tvm.tirx.Var("i", "int32"), tvm.tirx.Var("j", "int32"), tvm.tirx.Var("k", "int32")
+    var_dom = {
+        i: tvm.ir.Range(begin=0, end=16),
+        j: tvm.ir.Range(begin=0, end=16),
+        k: tvm.ir.Range(begin=0, end=16),
+    }
+    check_region_bound(
+        {i // 4: None, j * 4 + i % 4: None, tirx.truncdiv(k, 2): None},
+        var_dom,
+        predicate=j * 4 + i % 4 > 3,
+        mode="lowerbound",
+    )
+    check_region_bound(
+        {i // 4: (0, 4), j * 4 + i % 4: (4, 64), tirx.truncdiv(k, 2): (0, 8)},
+        var_dom,
+        predicate=j * 4 + i % 4 > 3,
+        mode="upperbound",
+    )
+
+
+def test_region_bound_stride_too_wide():
+    i = tvm.tirx.Var("i", "int32")
+    var_dom = {i: tvm.ir.Range(begin=0, end=64)}
+    check_region_bound({(i * 4, i * 4 + 2): None}, var_dom, mode="lowerbound")
+    check_region_bound({(i * 4, i * 4 + 2): (0, 254)}, var_dom, mode="upperbound")
+
+
+def test_region_bound_small_stride():
+    i = tvm.tirx.Var("i", "int32")
+    var_dom = {
+        i: tvm.ir.Range(begin=0, end=64),
+    }
+    check_region_bound({(i * 4, i * 4 + 8): (0, 260)}, var_dom, mode="lowerbound")
+
+
+def test_region_lower_bound_split_predicate():
+    x_o = tvm.tirx.Var("xo", "int32")
+    x_i = tvm.tirx.Var("xi", "int32")
+    x = x_o * 4 + x_i
+    var_dom = {
+        x_o: tvm.ir.Range(begin=0, end=16),
+        x_i: tvm.ir.Range(begin=0, end=4),
+    }
+    check_region_bound({(x * 4, x * 4 + 8): (0, 256)}, var_dom, predicate=x < 63, mode="lowerbound")
+
+    check_region_bound(
+        {(x * 4, x * 4 + 8): (0, 256), (x * 3, x * 3 + 5): (0, 191)},
+        var_dom,
+        predicate=x < 63,
+        mode="upperbound",
+    )
+
+
+def test_region_lower_bound_multiple_variables():
+    div = tvm.tirx.floordiv
+    mod = tvm.tirx.floormod
+    x = tvm.tirx.Var("x", "int32")
+    wid = tvm.tirx.Var("wid", "int32")
+    i = div(x, 16)
+    j = div(mod(x, 16), 4) * 8 + mod(x, 4) + div(wid, 32) * 4
+    k = wid % 32
+    var_dom = {
+        x: tvm.ir.Range(begin=0, end=32),
+        wid: tvm.ir.Range(begin=0, end=64),
+    }
+    check_region_bound({i: (0, 2), j: (0, 32), k: (0, 32)}, var_dom, mode="lowerbound")
+
+
+def test_region_lower_bound_negative_scale():
+    i = tvm.tirx.Var("i", "int32")
+    j = tvm.tirx.Var("j", "int32")
+    var_dom = {
+        i: tvm.ir.Range(begin=0, end=4),
+        j: tvm.ir.Range(begin=0, end=4),
+    }
+    check_region_bound(
+        {(1 - i, 5 - i): (-2, 5), (20 - j * 4, 36 - j * 4): (8, 36)}, var_dom, mode="lowerbound"
+    )
+
+
+def test_region_lower_bound_for_non_perfect_tile():
+    h1 = tvm.tirx.Var("h1", "int32")
+    h2 = tvm.tirx.Var("h2", "int32")
+    h3 = tvm.tirx.Var("h3", "int32")
+
+    # non-uniform tiling, single inner variable
+    var_dom = {
+        h2: tvm.ir.Range(begin=0, end=10),
+    }
+    check_region_bound(
+        {
+            h3 * 8 + h2: {
+                (): (
+                    tvm.tirx.max(h3 * 8, 1),
+                    tvm.tirx.min(0, h3 * 8 - 214) + 224,
+                ),
+                ((h3, 0),): (1, 10),  # h3 == 0: region is [1, 10)
+                ((h3, 10),): (h3 * 8, h3 * 8 + 10),  # 0 < h3 <= 26: region is [h3 * 8, h3 * 8 + 10)
+                ((h3, 27),): (h3 * 8, 224),  # h3 > 26: region is [h3 * 8, 224)
+            }
+        },
+        var_dom,
+        predicate=tvm.tirx.all(1 <= h3 * 8 + h2, h3 * 8 + h2 < 224),
+        mode="lowerbound",
+    )
+
+    # non-uniform tiling, two inner variables
+    var_dom = {
+        h1: tvm.ir.Range(begin=0, end=5),
+        h2: tvm.ir.Range(begin=0, end=2),
+    }
+    check_region_bound(
+        {
+            h3 * 8 + h2 * 5 + h1: {
+                (): (
+                    tvm.tirx.max(h3 * 8, 1),
+                    tvm.tirx.min(0, h3 * 8 - 214) + 224,
+                ),
+                ((h3, 0),): (1, 10),
+                ((h3, 10),): (h3 * 8, h3 * 8 + 10),
+                ((h3, 27),): (h3 * 8, 224),
+            }
+        },
+        var_dom,
+        predicate=tvm.tirx.all(1 <= h3 * 8 + h2 * 5 + h1, h3 * 8 + h2 * 5 + h1 < 224),
+        mode="lowerbound",
+    )
+
+    # lowerbound should fail on incompatible predicates
+    check_region_bound(
+        {h3 * 8 + h2 * 5 + h1: None},
+        var_dom,
+        predicate=tvm.tirx.all(1 <= h3 * 8 + h2 * 5 + h1, h3 * 8 + h1 * 2 + h2 < 224),
+        mode="lowerbound",
+    )
+    check_region_bound(
+        {h3 * 8 + h2 * 5 + h1: (h3 * 8, h3 * 8 + 10)},
+        var_dom,
+        predicate=tvm.tirx.all(1 <= h3 * 8 + h2 * 5 + h1, h3 * 8 + h1 * 2 + h2 < 224),
+        mode="upperbound",
+    )
+
+
+def test_region_lower_bound_unfusable():
+    var_dom = {
+        tvm.tirx.Var("i", "int32"): tvm.ir.Range(8),
+        tvm.tirx.Var("j", "int32"): tvm.ir.Range(4),
+    }
+    i, j = var_dom
+    check_region_bound({(i + j) // 2: (0, 6)}, var_dom, mode="lowerbound")
+
+
+def test_union_lower_bound():
+    neg_inf = tvm.arith.int_set.neg_inf()
+    pos_inf = tvm.arith.int_set.pos_inf()
+    set_0 = tvm.arith.IntervalSet(min_value=neg_inf, max_value=0)
+    set_1 = tvm.arith.IntervalSet(min_value=1, max_value=pos_inf)
+    result = tvm.arith.int_set.union_lower_bound([set_0, set_1])
+    assert result.min_value.same_as(neg_inf)
+    assert result.max_value.same_as(pos_inf)
+    set_2 = tvm.arith.IntervalSet(min_value=pos_inf, max_value=neg_inf)
+    result = tvm.arith.int_set.union_lower_bound([set_0, set_1, set_2])
+    assert result.min_value.same_as(neg_inf)
+    assert result.max_value.same_as(pos_inf)
+
+
+def test_modular_set():
+    ck = IntSetChecker()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    expr = (x * 2048 + y * 16) % 7168
+    ck.verify(
+        expr, {x: tvm.arith.IntervalSet(0, 128), y: tvm.arith.IntervalSet(0, 3584)}, (0, 7152)
+    )
+
+
+def test_relax_deep_variable_dependency_chain():
+    """Regression test for exponential variable-relaxation blowup.
+
+    When a variable's interval bound references another variable that is also in
+    the domain map, the evaluator relaxes it transitively. A diamond-shaped
+    chain -- where each variable's bound references the next one in *both* its
+    min and its max -- used to be re-expanded along every path, costing
+    O(2^depth) and hanging indefinitely. The relaxation is now memoized per
+    variable, so this completes in linear time.
+    """
+    ck = IntSetChecker()
+    n = 64  # 2^64 expansions without memoization; trivially fast with it.
+    xs = [tvm.tirx.Var(f"x{i}", "int32") for i in range(n + 1)]
+    dmap = {xs[i]: tvm.arith.IntervalSet(xs[i + 1] - 1, xs[i + 1] + 1) for i in range(n)}
+    dmap[xs[n]] = tvm.arith.IntervalSet(0, 100)
+    # x0 relaxes through the whole chain: [0 - n, 100 + n].
+    ck.verify(xs[0], dmap, (-n, 100 + n))
+
+
+def test_relax_cyclic_variable_dependency():
+    """A cyclic variable dependency must terminate (and stay symbolic)."""
+    ana = tvm.arith.Analyzer()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    # x depends on y and y depends on x: relaxation must not loop forever.
+    dmap = {x: tvm.arith.IntervalSet(y, y), y: tvm.arith.IntervalSet(x, x)}
+    res = ana.int_set(x, dmap)
+    assert res is not None
+
+
+def test_estimate_region_accepts_external_analyzer():
+    i = tvm.tirx.Var("i", "int32")
+    tile = tvm.tirx.Var("tile", "int32")
+    region = [tvm.ir.Range.from_min_extent(i % tile, 1)]
+    dom = {i: tvm.ir.Range(0, 16)}
+
+    # Without knowing `tile`, the affine detection fails for exact bounds.
+    assert tvm.arith.estimate_region_lower_bound(region, dom, True) is None
+    assert tvm.arith.estimate_region_strict_bound(region, dom, True) is None
+    upper_without_analyzer = tvm.arith.estimate_region_upper_bound(region, dom, True)
+
+    analyzer = tvm.arith.Analyzer()
+    analyzer.bind(tile, tvm.tirx.const(4, "int32"))
+    # The external binding lets the affine detection succeed.
+    for estimate_region in [
+        tvm.arith.estimate_region_lower_bound,
+        tvm.arith.estimate_region_strict_bound,
+        tvm.arith.estimate_region_upper_bound,
+    ]:
+        result = estimate_region(region, dom, True, analyzer=analyzer)
+        assert result is not None
+        assert analyzer.can_prove_equal(result[0].min_value, 0)
+        assert analyzer.can_prove_equal(result[0].max_value, 3)
+
+    # The upper-bound fallback without analyzer is safe but much wider.
+    assert not analyzer.can_prove_equal(upper_without_analyzer[0].min_value, 0)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_modular_set.py

@@ -0,0 +1,234 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+import tvm
+import tvm.testing
+from tvm import te
+
+
+def test_cast():
+    analyzer = tvm.arith.Analyzer()
+    x = tvm.tirx.Var("x", "int8")
+    m = analyzer.modular_set((x * 3).astype("uint32"))
+    assert m.coeff == 3
+    assert m.base == 0
+    m = analyzer.modular_set((x * 3 + 1).astype("float32").astype("int32"))
+    assert m.coeff == 3
+    assert m.base == 1
+
+
+def test_add_sub():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int64"), tvm.tirx.Var("y", "int64")
+    m = analyzer.modular_set(x * 6 + y * 4)
+    assert m.coeff == 2
+    assert m.base == 0
+
+    analyzer.bind(y, x * 4 + 1)
+    m = analyzer.modular_set(1 - y)
+    assert m.coeff == 4
+    assert m.base == 0
+
+
+def test_mul():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    m = analyzer.modular_set((x * 4 + 2) * (y * 6 + 1))
+    assert m.coeff == 4
+    assert m.base == 2
+
+
+def test_shift_left():
+    analyzer = tvm.arith.Analyzer()
+    x, y = te.var("x"), te.var("y")
+    m = analyzer.modular_set((x * 4 + 2) << 2)
+    assert m.coeff == 16
+    assert m.base == 8
+
+
+def test_floormod():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    m = analyzer.modular_set(tvm.tirx.floormod(x * 128 + y * 4, 256))
+    assert m.coeff == 4
+    assert m.base == 0
+
+
+def test_div_shift():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    # not sure if x is non-negative
+    tdiv = tvm.tirx.truncdiv
+    m = analyzer.modular_set(tdiv(x * 4 + 2, 2))
+    assert m.coeff == 1
+    assert m.base == 0
+    # right shift always round down so it is fine
+    m = analyzer.modular_set((x * 4 + 2) >> 1)
+    assert m.coeff == 2
+    assert m.base == 1
+    fld = tvm.tirx.floordiv
+    m = analyzer.modular_set(fld(x * 4 + 2, 2))
+    assert m.coeff == 2
+    assert m.base == 1
+    # x is non-negative
+    analyzer.update(x, tvm.arith.ConstIntBound(0, 100))
+    m = analyzer.modular_set(tdiv(x * 4 + 2, 2))
+    assert m.coeff == 2
+    assert m.base == 1
+
+
+def test_mod():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    tmod = tvm.tirx.truncmod
+    fmod = tvm.tirx.floormod
+    # not sure if x is non-negative
+    m = analyzer.modular_set(tmod(x * 4 + 1, 4))
+    assert m.coeff == 1
+    assert m.base == 0
+    # no need to be positive if base == 0
+    m = analyzer.modular_set(tmod(x * 4, 4))
+    assert m.coeff == 4
+    assert m.base == 0
+    # floor mod tests
+    m = analyzer.modular_set(fmod(x * 4 + 3, 2))
+    assert m.coeff == 2
+    assert m.base == 1
+    m = analyzer.modular_set(fmod(x * 4 + 3, 8))
+    assert m.coeff == 4
+    assert m.base == 3
+    # x is non-negative
+    analyzer.update(x, tvm.arith.ConstIntBound(0, 100))
+    m = analyzer.modular_set(tmod(x * 4 + 3, 2))
+    assert m.coeff == 2
+    assert m.base == 1
+
+
+def test_min_max_select():
+    analyzer = tvm.arith.Analyzer()
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    m = analyzer.modular_set(tvm.tirx.min(x * 3, y * 9))
+    assert m.coeff == 3
+    assert m.base == 0
+
+    m = analyzer.modular_set(tvm.tirx.max(x * 3 + 1, y * 9 + 4))
+    assert m.coeff == 3
+    assert m.base == 1
+
+    m = analyzer.modular_set(tvm.tirx.Select(x > 0, x * 3 + 1, y * 9 + 2))
+    assert m.coeff == 1
+    assert m.base == 0
+
+
+def test_mix_index():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    analyzer = tvm.arith.Analyzer()
+    tdiv = tvm.tirx.truncdiv
+    m = analyzer.modular_set(a * 4 + b * 6 + 7)
+    assert m.coeff == 2
+    assert m.base == 1
+
+    m = analyzer.modular_set((a * 4 + 1) * (b * 8 + 3))
+    assert m.coeff == 4
+    assert m.base == 3
+
+    m = analyzer.modular_set(tdiv(a * 4 + 1, b * 8 + 3))
+    assert m.coeff == 1
+    assert m.base == 0
+
+    m = analyzer.modular_set((a * 4 + 1) * tdiv(b * 8, 4))
+    assert m.coeff == 2
+    assert m.base == 0
+
+    m = analyzer.modular_set((a * 12 + 1) - (b * 3 * 7 + 2))
+    assert m.coeff == 3
+    assert m.base == 2
+
+    m = analyzer.modular_set(a * 12 + tvm.tirx.min(b * 3 * 7, 2))
+    assert m.coeff == 1
+    assert m.base == 0
+
+
+def test_constraint_scope():
+    a = tvm.tirx.Var("a", "int32")
+    b = tvm.tirx.Var("b", "int32")
+    analyzer = tvm.arith.Analyzer()
+    tmod = tvm.tirx.truncmod
+
+    with analyzer.constraint_scope(tmod(b, 4) == 2):
+        m = analyzer.modular_set(b + 1)
+        assert m.coeff == 4
+        assert m.base == 3
+        with analyzer.constraint_scope(tmod(a, 2) == 1):
+            m = analyzer.modular_set(b + a * 2)
+            assert m.coeff == 4
+            assert m.base == 0
+        m = analyzer.modular_set(b + a * 2)
+        assert m.coeff == 2
+        assert m.base == 0
+
+    m = analyzer.modular_set(b + 1)
+    assert m.coeff == 1
+    assert m.base == 0
+
+
+def test_intersect():
+    a = tvm.tirx.Var("a", "int32")
+    analyzer = tvm.arith.Analyzer()
+    tmod = tvm.tirx.truncmod
+    with analyzer.constraint_scope(tmod(a, 4) == 1):
+        with analyzer.constraint_scope(tmod(a, 3) == 1):
+            m = analyzer.modular_set(a)
+            assert m.coeff == 12
+            assert m.base == 1
+
+    with analyzer.constraint_scope(tmod(a, 3) == 2):
+        with analyzer.constraint_scope(tmod(a, 5) == 3):
+            with analyzer.constraint_scope(tmod(a, 7) == 2):
+                m = analyzer.modular_set(a)
+                assert m.coeff == 105
+                assert m.base == 23
+
+
+def test_let():
+    analyzer = tvm.arith.Analyzer()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+    m = analyzer.modular_set(tvm.tirx.Let(x, y * 10, x + 1))
+    assert m.coeff == 10
+    assert m.base == 1
+
+
+def test_bitwise_and():
+    analyzer = tvm.arith.Analyzer()
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "int32")
+
+    # RHS of bitwise_and is 2^p - 1
+    m = analyzer.modular_set((x * 16 + y * 4) & 31)
+    assert m.coeff == 4
+    assert m.base == 0
+
+    # arbitrary RHS
+    m = analyzer.modular_set((x * 16 + y * 4) & 17)
+    assert m.coeff == 1
+    assert m.base == 0
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_simplify.py

@@ -0,0 +1,195 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import tvm
+import tvm.ir
+import tvm.testing
+from tvm import tirx
+from tvm.script import tirx as T
+
+
+def test_simplify_reshape_flattened_index():
+    ana = tvm.arith.Analyzer()
+
+    i0 = tirx.Var("i0", "int64")
+    i1 = tirx.Var("i1", "int64")
+    ana.bind(i0, tvm.ir.Range(0, 8))
+    ana.bind(i1, tvm.ir.Range(0, 3))
+
+    i_flattened = i0 * 3 + i1
+    tvm.ir.assert_structural_equal(
+        ana.simplify((i_flattened) // 12 * 12 + (i_flattened) % 12 // 4 * 4 + (i_flattened) % 4),
+        i_flattened,
+    )
+
+
+dtype = tvm.testing.parameter(
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "float16",
+    "float32",
+    "float64",
+)
+
+
+def test_can_prove_self_identity(dtype):
+    ana = tvm.arith.Analyzer()
+
+    n = tirx.Var("n", dtype)
+    assert ana.can_prove(n == n)
+
+
+def test_can_prove_self_equal_to_self(dtype):
+    ana = tvm.arith.Analyzer()
+
+    n = tirx.Var("n", dtype)
+    assert ana.can_prove_equal(n, n)
+
+
+def test_simplify_symbolic_comparison():
+    ana = tvm.arith.Analyzer()
+
+    i0 = tirx.Var("i0", "int64")
+    i1 = tirx.Var("i1", "int64")
+    n, m = tvm.tirx.Var("n", "int64"), tvm.tirx.Var("m", "int64")
+    outer = (n + 31) // 32
+    PS = tvm.arith.ProofStrength
+
+    non_negative = tvm.arith.ConstIntBound(0, tvm.arith.ConstIntBound.POS_INF)
+    ana.update(n, non_negative)
+    ana.update(m, non_negative)
+    ana.bind(i0, tvm.ir.Range(0, outer))
+    ana.bind(i1, tvm.ir.Range(0, 32))
+    assert not ana.can_prove(i0 * 32 + i1 < (n + 31) // 32 * 32, PS.DEFAULT)
+    assert ana.can_prove(i0 * 32 + i1 < (n + 31) // 32 * 32, PS.SYMBOLIC_BOUND)
+    assert ana.can_prove(i0 * 32 + i1 < (n + 31) // 32 * 32 + m, PS.SYMBOLIC_BOUND)
+    assert ana.can_prove(i0 * 32 + i1 + 1 <= (n + 31) // 32 * 32, PS.SYMBOLIC_BOUND)
+    assert ana.can_prove((n + 31) // 32 * 32 >= i0 * 32 + i1 + 1, PS.SYMBOLIC_BOUND)
+    assert ana.can_prove((n + 31) // 32 * 32 >= i0 * 32 + i1, PS.SYMBOLIC_BOUND)
+
+
+# These tests exercised arith::CanProve's substitution-based proof loop for
+# vscale-bearing expressions (iterating over known vscale values for a VLA target).
+# That loop has been removed -- arith no longer attempts target-dependent proofs
+# about scalable-vector lengths. The LOG(WARNING) for non-VLA targets is also gone.
+@pytest.mark.xfail(reason="arith no longer proves vscale-bearing inequalities via substitution")
+@pytest.mark.parametrize(
+    "expression",
+    [
+        T.vscale() * 32 < T.vscale() * 64,
+        T.vscale() * 2 * (T.vscale() * 2) >= T.vscale() * 4,
+        (T.vscale() * 4 + 114) // (T.vscale() * 4) * (T.vscale() * 4) >= 115,
+        64 % T.vscale() <= T.vscale(),
+    ],
+)
+def test_simplify_vscale_comparison_with_sve_target(expression):
+    ana = tvm.arith.Analyzer()
+
+    with tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}):
+        assert ana.can_prove(expression)
+
+
+@pytest.mark.xfail(
+    reason="arith no longer emits a LOG(WARNING) for vscale proofs on non-VLA targets"
+)
+def test_simplify_vscale_comparison_without_sve_target(capfd):
+    ana = tvm.arith.Analyzer()
+    vs = tvm.tirx.vscale()
+
+    with pytest.raises(AssertionError):
+        with tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu"}):
+            assert ana.can_prove(vs * 32 < vs * 64)
+
+    warning_prefix = (
+        "Warning: The expression contains scalable values. An attempt to prove by substituting "
+        "with known values of vscale was not performed. This proof currently only supports "
+        "VLA targets, but the target was "
+    )
+    capture = capfd.readouterr().err
+    assert warning_prefix in capture
+    assert '"kind":"llvm"' in capture
+    assert '"mtriple":"aarch64-linux-gnu"' in capture
+
+
+def test_regression_simplify_inf_recursion():
+    ana = tvm.arith.Analyzer()
+    cond = tirx.Var("cond", "int32")
+
+    res = (tvm.tirx.NE(cond, 0).astype("int8") - tvm.tirx.NE(cond, 0).astype("int8")).astype(
+        "int32"
+    ) == 0
+    # regression in a previous case
+    # try compare and int set recursive call can cause infinite loop
+    ana.rewrite_simplify(res)
+
+
+def test_bind_allow_override():
+    ana = tvm.arith.Analyzer()
+    x = tirx.Var("x", "int64")
+
+    ana.bind(x, tvm.ir.Range(0, 10))
+    ana.bind(x, tvm.ir.Range(0, 5), allow_override=True)
+    assert ana.can_prove(x < 5)
+
+    with pytest.raises(RuntimeError, match="Trying to update var 'x' with a different const bound"):
+        ana.bind(x, tvm.ir.Range(0, 3))
+
+
+def test_simplify_floor_mod_with_linear_offset():
+    """
+    Test that the floor_mod is simplified correctly when the offset is linear.
+    """
+    ana = tvm.arith.Analyzer()
+    past_decoder_sequence_length = tirx.Var("past_decoder_sequence_length", "int64")
+    expr1 = (past_decoder_sequence_length + 1) * 64
+    divisor1 = (past_decoder_sequence_length + 1) * 32
+    assert ana.can_prove_equal(tvm.tirx.floormod(expr1, divisor1), 0)
+    divisor2 = 32 * (past_decoder_sequence_length + 1)
+    assert ana.can_prove_equal(tvm.tirx.floormod(expr1, divisor2), 0)
+
+
+def test_simplify_uint_floormod_const_scale_divisible():
+    """uint32 floormod(x * c1, c2) -> 0 when c1 % c2 == 0 (overflow-free)."""
+    ana = tvm.arith.Analyzer()
+    q = tirx.Var("q_stage_idx", "uint32")
+    expr = q * tirx.Cast("uint32", 128)
+    mod = expr % tirx.const(4, "uint32")
+    assert ana.can_prove_equal(mod, tirx.const(0, "uint32"))
+    tvm.ir.assert_structural_equal(ana.rewrite_simplify(mod), tirx.const(0, "uint32"))
+
+
+def test_simplify_float_division():
+    # Test for the discussion:
+    # https://discuss.tvm.apache.org/t/discuss-is-constant-division-to-multiplication-rewrite-in-tvm-necessary/18615
+    ana = tvm.arith.Analyzer()
+    x = tirx.Var("x", "float32")
+    ry = x / 27
+    # in old version, the division will be rewritten into x * T.float32(1 / 27)
+    sy = ana.rewrite_simplify(ry)
+    tvm.ir.assert_structural_equal(ry, sy)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_solve_linear_equations.py

@@ -0,0 +1,186 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+import random
+import sys
+
+import pytest
+import tvm_ffi
+
+import tvm
+from tvm import arith, ir, testing, tirx
+from tvm.script import tirx as T
+
+
+def test_solution_consistency():
+    seed = random.randrange(sys.maxsize)
+    print(
+        "\nThis test is intentionally non-deterministic, "
+        f"if it fails please report it in GitHub issue together with this seed {seed}\n"
+    )
+    random.seed(seed)
+
+    def _check(num_vars, num_formulas, coef=(-5, 5), bounds=(-20, 20)):
+        variables = [tvm.tirx.Var("x" + str(i), "int32") for i in range(num_vars)]
+
+        relations = []
+        for i in range(num_formulas):
+            s1 = sum([v * random.randint(coef[0], coef[1]) for v in variables])
+            s1 += random.randint(coef[0], coef[1])
+            s2 = sum([v * random.randint(coef[0], coef[1]) for v in variables])
+            s2 += random.randint(coef[0], coef[1])
+            if random.random() < 0.7:
+                op = tvm.tirx.EQ
+            else:
+                # we also make sure it can correctly handle inequalities
+                op = random.choice([tvm.tirx.LE, tvm.tirx.LT, tvm.tirx.GE, tvm.tirx.GT])
+            relations.append(op(s1, s2))
+
+        vranges = {v: tvm.ir.expr.Range(bounds[0], bounds[1] + 1) for v in variables}
+        solution = arith.solve_linear_equations(relations, variables, vranges)
+
+        testing.check_int_constraints_trans_consistency(solution)
+
+        # leaving some variables as parameters should also be ok
+        for k in [1, 2]:
+            if len(variables) > k:
+                solution = arith.solve_linear_equations(relations, variables[:-k], vranges)
+                param_ranges = {v: vranges[v] for v in variables[-k:]}
+                testing.check_int_constraints_trans_consistency(solution, param_ranges)
+
+    for i in range(2):
+        _check(num_vars=1, num_formulas=1)
+    for i in range(2):
+        _check(num_vars=1, num_formulas=2)
+
+    for i in range(2):
+        _check(num_vars=2, num_formulas=1)
+    for i in range(2):
+        _check(num_vars=2, num_formulas=2)
+    for i in range(2):
+        _check(num_vars=2, num_formulas=3)
+
+    for i in range(3):
+        _check(num_vars=3, num_formulas=3, coef=(-2, 2))
+    for i in range(3):
+        _check(num_vars=3, num_formulas=4, coef=(-2, 2))
+
+    for i in range(3):
+        _check(num_vars=4, num_formulas=3, coef=(-1, 1))
+
+    for i in range(3):
+        _check(num_vars=10, num_formulas=2, coef=(-1, 1), bounds=(0, 4))
+    for i in range(3):
+        _check(num_vars=10, num_formulas=3, coef=(0, 1), bounds=(0, 4))
+
+
+def test_empty_var_to_solve():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    equations = [
+        tvm.tirx.EQ(x + y, 20),
+        tvm.tirx.EQ(x - y, 10),
+    ]
+    solution = arith.solve_linear_equations(equations)
+    assert len(solution.src_to_dst) == 0
+    assert len(solution.dst_to_src) == 0
+    assert len(solution.src.variables) == 0
+    assert len(solution.src.ranges) == 0
+    assert tvm_ffi.structural_equal(solution.src.relations, equations)
+    assert tvm_ffi.structural_equal(solution.src, solution.dst)
+
+
+def test_unique_solution():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    solution = arith.solve_linear_equations(
+        [
+            tvm.tirx.EQ(x + y, 20),
+            tvm.tirx.EQ(x - y, 10),
+        ],
+        [x, y],
+    )
+    assert list(solution.dst.variables) == []
+    assert tvm_ffi.structural_equal(solution.src_to_dst[x], T.int32(15))
+    assert tvm_ffi.structural_equal(solution.src_to_dst[y], T.int32(5))
+
+
+def test_low_rank():
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+    ranges = {}
+
+    solution = arith.solve_linear_equations(
+        [
+            tvm.tirx.EQ(x + y + z, 15),
+            tvm.tirx.EQ(x + y, 10),
+        ],
+        [x, y, z],
+        ranges,
+    )
+    [n0] = solution.dst.variables
+    assert tvm_ffi.structural_equal(solution.src_to_dst[x], n0 + 10)
+    assert tvm_ffi.structural_equal(solution.src_to_dst[y], -n0)
+    assert tvm_ffi.structural_equal(solution.src_to_dst[z], T.int32(5))
+
+
+def test_infer_range():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    ranges = {
+        x: tvm.ir.Range.from_min_extent(-5, 10),
+        y: tvm.ir.Range.from_min_extent(0, 10),
+    }
+
+    solution = arith.solve_linear_equations(
+        [
+            tvm.tirx.EQ(x + y, 0),
+        ],
+        [x, y],
+        ranges,
+    )
+    [n0] = solution.dst.variables
+    assert tvm_ffi.structural_equal(solution.src_to_dst[x], n0)
+    assert tvm_ffi.structural_equal(solution.src_to_dst[y], -n0)
+    # inferred from y's range
+    assert tvm_ffi.structural_equal(solution.dst.ranges[n0].min, T.int32(-9))
+    assert tvm_ffi.structural_equal(solution.dst.ranges[n0].extent, T.int32(10))
+    # additional inequality is added into the system for x
+    [ineq] = solution.dst.relations
+    assert isinstance(ineq, tvm.tirx.LE)
+    assert tvm_ffi.structural_equal(ineq.a, T.int32(-5))
+    assert tvm_ffi.structural_equal(ineq.b, n0)
+
+
+def test_ill_formed():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    solution = arith.solve_linear_equations(
+        [
+            tvm.tirx.EQ(x + y, 0),
+            tvm.tirx.EQ(x - y, 0),
+            tvm.tirx.EQ(x, 5),
+        ],
+        [x, y],
+        {},
+    )
+    assert list(solution.dst.variables) == []
+    [rel] = solution.dst.relations
+    ir.assert_structural_equal(rel, tirx.const(False))
+    assert len(solution.src_to_dst) == 0
+    assert len(solution.dst_to_src) == 0
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_solve_linear_inequality.py

@@ -0,0 +1,223 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import random
+import sys
+
+import pytest
+import tvm_ffi
+
+import tvm
+from tvm import arith, ir, testing, tirx
+from tvm.script import tirx as T
+
+
+@pytest.mark.skip(reason="See https://github.com/apache/tvm/issues/11458")
+def test_solution_consistency():
+    seed = random.randrange(sys.maxsize)
+    print(
+        "\nThis test is intentionally non-deterministic, "
+        f"if it fails please report it in GitHub issue together with this seed {seed}\n"
+    )
+    random.seed(seed)
+
+    def _check(variables, formulas, coef=(-5, 5), bounds=(-20, 20)):
+        vs = [tvm.tirx.Var("x" + str(i), "int32") for i in range(variables)]
+
+        fs = []
+        for i in range(formulas):
+            s1 = sum([v * random.randint(coef[0], coef[1]) for v in vs])
+            s1 += random.randint(coef[0], coef[1])
+            s2 = sum([v * random.randint(coef[0], coef[1]) for v in vs])
+            s2 += random.randint(coef[0], coef[1])
+            op = random.choice(
+                [tirx.expr.EQ, tirx.expr.LE, tirx.expr.LT, tirx.expr.GE, tirx.expr.GT]
+            )
+            fs.append(op(s1, s2))
+
+        vranges = {v: tvm.ir.expr.Range(bounds[0], bounds[1] + 1) for v in vs}
+        before = tvm.tirx.all(tirx.const(1, "bool"), *fs)
+        after = arith._ffi_api.SolveInequalitiesAsCondition(vs, vranges, fs)
+        after = tvm.tirx.all(tirx.const(1, "bool"), *after)
+        testing.check_bool_expr_is_true(before == after, vranges)
+
+        solution = arith.solve_linear_inequalities(fs, vs, vranges, deskew_range=True)
+        testing.check_int_constraints_trans_consistency(solution)
+
+    for i in range(3):
+        _check(1, 1)
+    for i in range(3):
+        _check(1, 2)
+
+    for i in range(3):
+        _check(2, 1)
+    for i in range(3):
+        _check(2, 2)
+    for i in range(3):
+        _check(2, 3)
+
+    # Somewhere here coefficients in the results become too large, leading to overflow,
+    # so we use smaller initial coefficients
+    for i in range(5):
+        _check(3, 3, coef=(-2, 2))
+    for i in range(5):
+        _check(3, 4, coef=(-2, 2))
+
+    for i in range(5):
+        _check(4, 3, coef=(-1, 1))
+
+    for i in range(5):
+        _check(10, 2, coef=(-1, 1), bounds=(0, 4))
+    for i in range(5):
+        _check(10, 3, coef=(0, 1), bounds=(0, 4))
+
+
+def test_dual_variable():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+
+    variables = [x, y]
+    ranges = {
+        x: tvm.ir.Range(-100, 100),
+        y: tvm.ir.Range(0, 10),
+    }
+    problem = [
+        tvm.tirx.LE(x + y, 20),
+        tvm.tirx.GE(x - y, 10),
+    ]
+
+    # solution as conditions
+    solution = arith._ffi_api.SolveInequalitiesAsCondition(variables, ranges, problem)
+    assert tvm_ffi.structural_equal(solution[0], x >= (y + 10))
+    assert tvm_ffi.structural_equal(solution[1], x <= (20 - y))
+    assert tvm_ffi.structural_equal(solution[2], y >= 0)
+    assert tvm_ffi.structural_equal(solution[3], y <= 5)
+
+    # solve and get the ranges
+    solution = arith.solve_linear_inequalities(problem, variables, ranges)
+    # 0 <= y <=5
+    assert solution.ranges[y].min == 0
+    assert solution.ranges[y].extent == 6
+    # y + 10 <= x <= 20 - y
+    assert tvm_ffi.structural_equal(solution.ranges[x].min, y + 10)
+    assert solution.ranges[x].extent == 11  # max(10 - 2y)
+
+    # deskew the solved ranges to be starting from zero
+    solution = arith.solve_linear_inequalities(problem, variables, ranges, deskew_range=True)
+    [x_new, y_new] = solution.dst.variables
+    [rel] = solution.dst.relations
+    assert tvm_ffi.structural_equal(rel, (y_new * 2) + x_new <= 10)
+    assert tvm_ffi.structural_equal(solution.dst.ranges[x_new].min, T.int32(0))
+    assert tvm_ffi.structural_equal(solution.dst.ranges[x_new].extent, T.int32(11))
+    assert tvm_ffi.structural_equal(solution.dst.ranges[y_new].min, T.int32(0))
+    assert tvm_ffi.structural_equal(solution.dst.ranges[y_new].extent, T.int32(6))
+    assert tvm_ffi.structural_equal(solution.src_to_dst[x], x_new + (y_new + 10))
+    assert tvm_ffi.structural_equal(solution.src_to_dst[y], y_new)
+    assert tvm_ffi.structural_equal(solution.dst_to_src[x_new], x - y - 10)
+    assert tvm_ffi.structural_equal(solution.dst_to_src[y_new], y)
+
+
+def test_equal():
+    x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
+    problem = [
+        tvm.tirx.GE(x + y, 10),
+        tvm.tirx.GE(x - y, 2),
+        tvm.tirx.LE(x, 6),
+    ]
+
+    solution = arith.solve_linear_inequalities(problem, [x, y])
+    assert solution.ranges[x].min == 6
+    assert solution.ranges[x].extent == 1
+    assert solution.ranges[y].min == 4
+    assert solution.ranges[y].extent == 1
+
+    solution = arith.solve_linear_inequalities(problem, [x, y], deskew_range=True)
+    assert len(solution.dst.variables) == 0
+    assert len(solution.dst.ranges) == 0
+    assert len(solution.dst.relations) == 0
+    assert solution.src_to_dst[x] == 6
+    assert solution.src_to_dst[y] == 4
+
+
+def test_multi_equal():
+    x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
+    problem = [
+        tvm.tirx.LE(x, 6),
+        tvm.tirx.GE(x, 6),
+        tvm.tirx.GE(x - z * y, 0),
+        tvm.tirx.LE(x - z * y, 0),
+    ]
+
+    solution = arith.solve_linear_inequalities(problem, [x, y, z])
+    assert solution.ranges[x].min == 6
+    assert solution.ranges[x].extent == 1
+    assert len(solution.relations) == 3
+    assert tvm_ffi.structural_equal(solution.relations[0], x == z * y)
+
+    assert isinstance(solution.relations[1], tvm.tirx.LE)
+    assert solution.relations[1].b == 0
+    assert isinstance(solution.relations[2], tvm.tirx.LE)
+    assert solution.relations[2].b == 0
+    # (z*y - 6) <= 0 && (6 - z*y) <= 0
+    ana = tvm.arith.Analyzer()
+    assert ana.simplify(solution.relations[1].a + solution.relations[2].a) == 0
+    assert tvm_ffi.structural_equal(
+        solution.relations[1].a, (z * y - 6)
+    ) or tvm_ffi.structural_equal(solution.relations[2].a, (z * y - 6))
+
+    solution = arith.solve_linear_inequalities(problem, [x, y, z], deskew_range=True)
+    assert solution.src_to_dst[y] == y
+    assert solution.src_to_dst[z] == z
+    assert solution.src_to_dst[x] == 6
+
+
+def test_no_solution():
+    x = tvm.tirx.Var("x0", "int32")
+    vranges = {x: tvm.ir.Range.from_min_extent(-20, 41)}
+    problem = [-x - 4 <= -5 * x + 2, x * 4 + 5 <= x * 5]
+
+    solution = arith.solve_linear_inequalities(problem, [x], vranges, deskew_range=True)
+    assert list(solution.dst.variables) == []
+    [rel] = solution.dst.relations
+    ir.assert_structural_equal(rel, tirx.const(False))
+    assert len(solution.src_to_dst) == 0
+    assert len(solution.dst_to_src) == 0
+
+    solution = arith.solve_linear_inequalities(problem, [x], vranges)
+    assert len(solution.variables) == 0
+    assert len(solution.ranges) == 0
+    [rel] = solution.relations
+    assert not rel
+
+
+def test_unbound_var_range():
+    x = tvm.tirx.Var("x0", "int32")
+    free_var = tvm.tirx.Var("fv", "int32")
+    vranges = {
+        x: tvm.ir.Range.from_min_extent(0, tvm.tirx.Cast("int32", 1 + tvm.tirx.log(free_var)))
+    }
+    problem = [x > 3]
+    solution = arith.solve_linear_inequalities(
+        problem,
+        [x],
+        vranges,
+    )
+    assert len(solution.variables) == 1
+    assert len(solution.ranges) == 0
+    assert len(solution.relations) == 3
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/arith/test_arith_z3.py

@@ -0,0 +1,756 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import gc
+import queue
+import threading
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import tirx
+from tvm.arith import Analyzer, ProofStrength
+
+# The Z3 prover is only consulted at the kSymbolicBound strength so the common
+# default path never pays the prover cost.
+SB = ProofStrength.SYMBOLIC_BOUND
+
+
+def _require_z3(analyzer):
+    if not analyzer.is_z3_enabled:
+        pytest.skip("Z3 prover is disabled in this build")
+
+
+def implies(x, y):
+    return tirx.Or(tirx.Not(x), y)
+
+
+# ---------------------------------------------------------------------------
+# API availability (works regardless of whether Z3 is built)
+# ---------------------------------------------------------------------------
+
+
+def test_z3_capability_query():
+    # `is_z3_enabled` is the supported way to detect the build configuration.
+    # The Z3-specific debug/config methods work only when it is True, and raise
+    # a clear error otherwise.
+    analyzer = Analyzer()
+    assert isinstance(analyzer.is_z3_enabled, bool)
+
+    if analyzer.is_z3_enabled:
+        assert isinstance(analyzer.get_smtlib2(), str)
+        assert isinstance(analyzer.get_z3_stats(), str)
+    else:
+        with pytest.raises(RuntimeError):
+            analyzer.get_smtlib2()
+        with pytest.raises(RuntimeError):
+            analyzer.get_z3_stats()
+        with pytest.raises(RuntimeError):
+            analyzer.set_z3_timeout_ms(1000)
+        with pytest.raises(RuntimeError):
+            analyzer.set_z3_rlimit(0)
+
+
+def test_z3_context_lifetime_outlives_worker_thread():
+    _require_z3(Analyzer())
+
+    result_queue = queue.Queue()
+
+    def worker():
+        try:
+            analyzer = Analyzer()
+            x = tirx.Var("x", "int32")
+            analyzer.bind(x, tvm.ir.Range(0, 16))
+            assert analyzer.can_prove(x >= 0, SB)
+            result_queue.put(("analyzer", analyzer))
+        except BaseException as err:  # pylint: disable=broad-exception-caught
+            result_queue.put(("error", err))
+
+    thread = threading.Thread(target=worker)
+    thread.start()
+    thread.join()
+
+    kind, payload = result_queue.get_nowait()
+    if kind == "error":
+        raise payload
+
+    del payload
+    gc.collect()
+
+
+# ---------------------------------------------------------------------------
+# Examples the native analyzer cannot prove but Z3 can.
+#
+# Each case asserts both that the native analyzers (kDefault, Z3 gated off)
+# fail and that Z3 (kSymbolicBound) succeeds. This demonstrates the added value
+# of the Z3 backend and that it is correctly gated behind kSymbolicBound.
+# ---------------------------------------------------------------------------
+
+
+def test_z3_floor_division_identity_constraint():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+
+    expr = ((b - a) // c) * c + a <= b
+    with analyzer.constraint_scope(tirx.all(a > 0, b > 0, c > 0)):
+        assert not analyzer.can_prove(expr)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floor_division_identity_via_bind_range():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+
+    analyzer.bind(a, tvm.ir.Range(1, 100000))
+    analyzer.bind(b, tvm.ir.Range(1, 100000))
+    analyzer.bind(c, tvm.ir.Range(1, 100000))
+
+    expr = ((b - a) // c) * c + a <= b
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_multiplication_monotonicity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    d = tirx.Var("d", "int32")
+
+    expr = implies(tirx.all(a < b, b < c, a * d < b * d), b * d < c * d)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_nested_floor_division_collapse():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    expr = implies(
+        tirx.all(a >= 0, a < 128),
+        a // 128 == (a // 64 * 32 + a % 32 // 16 * 8) // 64,
+    )
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_deeply_nested_floor_division_identity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    expr = implies(
+        tirx.all(a >= 0, a < 128),
+        (
+            a % 16 * 64
+            + a // 64 * 32
+            + a % 8 // 4 * 32
+            + (a % 32 // 16 + a % 2) % 2 * 8
+            + 16
+            - (a // 64 + a % 8 // 4) // 2 * 64
+        )
+        // 512
+        == (
+            a % 16 * 64
+            + a // 64 * 32
+            + a % 8 // 4 * 32
+            + (a % 32 // 16 + a % 2) % 2 * 8
+            - (a // 64 + a % 8 // 4) // 2 * 64
+        )
+        // 512,
+    )
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_min_max_sum_identity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    expr = tirx.max(x, y) + tirx.min(x, y) == x + y
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_select_absolute_value_nonneg():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    expr = tirx.Select(x >= 0, x, -x) >= 0
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_transitive_inequality():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    expr = implies(tirx.all(a <= b, b <= c), a <= c)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_square_expansion_nonneg():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = (a + b) * (a + b) >= a * a + b * b
+    with analyzer.constraint_scope(tirx.all(a >= 0, b >= 0)):
+        assert not analyzer.can_prove(expr)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_square_monotonicity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = implies(tirx.all(0 <= a, a <= b), a * a <= b * b)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_strict_multiplication():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    d = tirx.Var("d", "int32")
+    expr = implies(tirx.all(a < b, d > 0), a * d < b * d)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floor_division_monotonicity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    expr = implies(tirx.all(a <= b, c > 0), tirx.floordiv(a, c) <= tirx.floordiv(b, c))
+    assert not analyzer.can_prove(expr)
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floor_division_lower_bound():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = implies(b > 0, tirx.floordiv(a, b) * b <= a)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floor_modulo_range():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = implies(b > 0, tirx.all(0 <= tirx.floormod(a, b), tirx.floormod(a, b) < b))
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_flattened_index_bound():
+    # Classic index-flattening bound used throughout TVM: for a row index i in
+    # [0, m) and a column index j in [0, n), the flattened index i * n + j stays
+    # within [0, m * n).
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    i = tirx.Var("i", "int32")
+    j = tirx.Var("j", "int32")
+    m = tirx.Var("m", "int32")
+    n = tirx.Var("n", "int32")
+    expr = tirx.all(0 <= i * n + j, i * n + j < m * n)
+    with analyzer.constraint_scope(tirx.all(0 <= i, i < m, 0 <= j, j < n, m > 0, n > 0)):
+        assert not analyzer.can_prove(expr)
+        analyzer.set_z3_rlimit(0)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_modular_combination():
+    # Native modular_set tracks single-variable moduli, but combining two
+    # independent modular facts to reason about their sum is left to Z3.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    expr = tirx.floormod(x + y, 2) == 0
+    with analyzer.constraint_scope(tirx.all(tirx.floormod(x, 6) == 0, tirx.floormod(y, 6) == 0)):
+        assert not analyzer.can_prove(expr)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_square_non_negative():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    assert not analyzer.can_prove(a * a >= 0)
+    assert analyzer.can_prove(a * a >= 0, SB)
+
+
+def test_z3_min_max_average_bounds():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    assert not analyzer.can_prove(tirx.max(a, b) * 2 >= a + b)
+    assert analyzer.can_prove(tirx.max(a, b) * 2 >= a + b, SB)
+    assert analyzer.can_prove(tirx.min(a, b) * 2 <= a + b, SB)
+
+
+def test_z3_symbolic_bind_range_with_constraint():
+    # Combine a symbolic range binding (x in [0, n)) with a constraint on the
+    # extent to derive a concrete bound on x.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    n = tirx.Var("n", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, n))
+    with analyzer.constraint_scope(n <= 8):
+        assert not analyzer.can_prove(x < 8)
+        assert analyzer.can_prove(x < 8, SB)
+
+
+def test_z3_equality_congruence():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = implies(a == b, a * a == b * b)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_integer_strict_transitivity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    # Over the integers, a < b and b < c implies a + 1 < c.
+    expr = implies(tirx.all(a < b, b < c), a + 1 < c)
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_if_then_else_absolute_value():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    expr = tirx.if_then_else(x >= 0, x, -x) >= 0
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_unsigned_non_negative():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    u = tirx.Var("u", "uint32")
+    assert not analyzer.can_prove(u >= 0)
+    assert analyzer.can_prove(u >= 0, SB)
+
+
+def test_z3_unsigned64_non_negative():
+    # Exercises the special-cased uint64 range handling (UINT64_MAX bound).
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    u = tirx.Var("u", "uint64")
+    assert not analyzer.can_prove(u >= 0)
+    assert analyzer.can_prove(u >= 0, SB)
+
+
+def test_z3_int64_square_expansion():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int64")
+    b = tirx.Var("b", "int64")
+    expr = (a + b) * (a + b) >= a * a + b * b
+    with analyzer.constraint_scope(tirx.all(a >= 0, b >= 0)):
+        assert not analyzer.can_prove(expr)
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_boolean_variable_reasoning():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    p = tirx.Var("p", "bool")
+    q = tirx.Var("q", "bool")
+    expr = implies(tirx.And(p, q), tirx.Or(p, q))
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_not_equal_from_strict_less():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    expr = implies(x < y, tirx.NE(x, y))
+    assert not analyzer.can_prove(expr)
+    assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_let_expression():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    y = tirx.Var("y", "int32")
+    t = tirx.Var("t", "int32")
+    let = tirx.Let(t, y * 2, t)
+    assert not analyzer.can_prove(let == y * 2)
+    assert analyzer.can_prove(let == y * 2, SB)
+
+
+def test_z3_cast_preserves_bounds():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    s = tirx.Var("s", "int16")
+    widened = tirx.Cast("int32", s)
+    assert analyzer.can_prove(widened <= 32767, SB)
+    assert analyzer.can_prove(widened >= -32768, SB)
+
+
+def test_z3_bitwise_and_mask_bound():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    assert analyzer.can_prove(tirx.bitwise_and(x, tirx.IntImm("int32", 7)) < 8, SB)
+
+
+def test_z3_bitwise_and_le_operand():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.bind(y, tvm.ir.Range(0, 256))
+    # Bit-vector reasoning over two variables exceeds the default deterministic
+    # rlimit; lift it (0 == unlimited, still deterministic) for this proof.
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(tirx.bitwise_and(x, y) <= x, SB)
+
+
+def test_z3_bitwise_or_ge_operand():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.bind(y, tvm.ir.Range(0, 256))
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(tirx.bitwise_or(x, y) >= x, SB)
+
+
+def test_z3_bitwise_xor_bound():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    y = tirx.Var("y", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.bind(y, tvm.ir.Range(0, 256))
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(tirx.bitwise_xor(x, y) < 256, SB)
+
+
+def test_z3_bitwise_not_identity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.set_z3_rlimit(0)
+    # Two's complement: ~x == -x - 1.
+    assert analyzer.can_prove(tirx.bitwise_not(x) == -x - 1, SB)
+
+
+def test_z3_shift_right_halves():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    analyzer.bind(x, tvm.ir.Range(0, 256))
+    analyzer.set_z3_rlimit(0)
+    # For non-negative x, (x >> 1) * 2 <= x.
+    assert analyzer.can_prove(tirx.shift_right(x, tirx.IntImm("int32", 1)) * 2 <= x, SB)
+
+
+def test_z3_shift_left_lower_bound():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    n = tirx.Var("n", "int32")
+    # Keep operands small so the 32-bit left shift cannot overflow; then
+    # x << n == x * 2 ** n >= x for x >= 1.
+    analyzer.bind(x, tvm.ir.Range(1, 16))
+    analyzer.bind(n, tvm.ir.Range(0, 4))
+    # Bit-vector shift reasoning exceeds the default deterministic rlimit.
+    analyzer.set_z3_rlimit(0)
+    assert analyzer.can_prove(tirx.shift_left(x, n) >= x, SB)
+
+
+# ---------------------------------------------------------------------------
+# Soundness / negative tests (Z3 must NOT prove false predicates)
+# ---------------------------------------------------------------------------
+
+
+def test_z3_negative_unprovable_inequality():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    # a < b does not hold for arbitrary a, b.
+    assert not analyzer.can_prove(a < b, SB)
+    # a * a > a is false (e.g. a == 0).
+    assert not analyzer.can_prove(a * a > a, SB)
+
+
+def test_z3_truncmod_can_be_negative():
+    # Regression test for truncated div/mod semantics: TVM Div/Mod round toward
+    # zero, so truncmod(a, 4) can be negative. A solver that modeled them as
+    # Euclidean would unsoundly "prove" truncmod(a, 4) >= 0.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    assert not analyzer.can_prove(tirx.truncmod(a, 4) >= 0, SB)
+
+
+def test_z3_truncdiv_truncmod_identity():
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    expr = tirx.truncdiv(a, b) * b + tirx.truncmod(a, b) == a
+    with analyzer.constraint_scope(b != 0):
+        assert analyzer.can_prove(expr, SB)
+
+
+def test_z3_floormod_nested_identities():
+    # Ported from TileLang's test_divmod. Here `%` is floormod: nested floormod
+    # by opposite-sign divisors collapses to the single-divisor result, while
+    # the mixed case does not.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    assert not analyzer.can_prove(a % 2 % -2 - a % 2 == 0, SB)
+    assert analyzer.can_prove(a % -2 % 2 - a % 2 == 0, SB)
+
+
+def test_z3_floormod_nonnegative():
+    # In contrast to truncmod, floormod with a positive divisor is always in
+    # [0, divisor), which Z3 should be able to prove.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    assert analyzer.can_prove(tirx.floormod(a, 4) >= 0, SB)
+    assert analyzer.can_prove(tirx.floormod(a, 4) < 4, SB)
+
+
+def test_z3_shift_does_not_poison_solver():
+    # Regression test: evaluating a shift expression must not add permanent
+    # assertions (such as `b >= 0` / `b < 64`) to the shared solver. Otherwise
+    # an unrelated, unbounded `b` would be wrongly provable to be < 100.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+
+    # Touch a shift expression so the prover visits the shift amount `b`.
+    analyzer.can_prove(tirx.shift_left(a, b) >= 0, SB)
+
+    # `b` is otherwise unconstrained, so this must remain unprovable.
+    assert not analyzer.can_prove(b < 100, SB)
+    assert not analyzer.can_prove(b >= 0, SB)
+
+
+def test_z3_constraint_scope_is_popped():
+    # Constraints entered through a scope must be removed once the scope exits,
+    # i.e. EnterConstraint's solver.push()/pop() must be balanced.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    with analyzer.constraint_scope(x > 5):
+        assert analyzer.can_prove(x > 0, SB)
+    # The constraint is gone; x is unconstrained again.
+    assert not analyzer.can_prove(x > 0, SB)
+
+
+def test_z3_opaque_call_is_safe():
+    # An opaque/unsupported sub-expression is modeled as a fresh free variable.
+    # It must neither crash nor be provable on its own, yet still be usable as a
+    # constraint.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    call = tirx.call_extern("int32", "foo", x)
+    assert not analyzer.can_prove(call > 0, SB)
+    with analyzer.constraint_scope(call > 0):
+        assert analyzer.can_prove(call > 0, SB)
+    assert not analyzer.can_prove(call > 0, SB)
+
+
+def test_z3_shift_overflow_is_not_proven():
+    # Z3 models fixed-width shifts via bit-vectors, so it correctly refuses to
+    # prove `x << n >= x` for an unbounded `x` (a large `x` overflows int32 and
+    # wraps to a negative value). This guards against unsound shift modeling.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    n = tirx.Var("n", "int32")
+    analyzer.set_z3_rlimit(0)
+    expr = implies(tirx.all(x >= 1, n >= 0, n < 8), tirx.shift_left(x, n) >= x)
+    assert not analyzer.can_prove(expr, SB)
+
+
+def test_z3_analyzers_are_isolated():
+    # Analyzers share a thread-local Z3 context but own separate solvers, so
+    # constraints and bindings in one must never leak into another.
+    analyzer_a = Analyzer()
+    analyzer_b = Analyzer()
+    _require_z3(analyzer_a)
+
+    x = tirx.Var("x", "int32")
+    with analyzer_a.constraint_scope(x > 100):
+        assert analyzer_a.can_prove(x > 50, SB)
+        assert not analyzer_b.can_prove(x > 50, SB)
+
+    analyzer_c = Analyzer()
+    analyzer_d = Analyzer()
+    analyzer_c.bind(x, tvm.ir.Range(0, 10))
+    assert analyzer_c.can_prove(x < 10, SB)
+    assert not analyzer_d.can_prove(x < 10, SB)
+
+
+def test_z3_repeated_can_prove_is_consistent():
+    # Repeated queries must be stateless: a CanProve call must not pollute the
+    # solver and change the result of a subsequent call.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    x = tirx.Var("x", "int32")
+    assert analyzer.can_prove(x > 0, SB) == analyzer.can_prove(x > 0, SB)
+
+    analyzer.bind(x, tvm.ir.Range(5, 10))
+    assert analyzer.can_prove(x >= 5, SB)
+    assert analyzer.can_prove(x >= 5, SB)
+
+
+def test_z3_is_gated_behind_symbolic_bound():
+    # The Z3 fallback must not run at the default strength.
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    expr = ((b - a) // c) * c + a <= b
+    with analyzer.constraint_scope(tirx.all(a > 0, b > 0, c > 0)):
+        assert not analyzer.can_prove(expr, ProofStrength.DEFAULT)
+        assert analyzer.can_prove(expr, SB)
+
+
+# ---------------------------------------------------------------------------
+# SMT-LIB2 export
+# ---------------------------------------------------------------------------
+
+
+def test_z3_smtlib2_roundtrip():
+    z3 = pytest.importorskip("z3")
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    expr = ((b - a) // c) * c + a <= b
+
+    solver = z3.Solver()
+    with analyzer.constraint_scope(tirx.all(a > 0, b > 0, c > 0)):
+        solver.from_string(analyzer.get_smtlib2(expr))
+    assert solver.check() == z3.unsat
+
+
+def test_z3_smtlib2_roundtrip_with_timeout():
+    z3 = pytest.importorskip("z3")
+    analyzer = Analyzer()
+    _require_z3(analyzer)
+
+    a = tirx.Var("a", "int32")
+    b = tirx.Var("b", "int32")
+    c = tirx.Var("c", "int32")
+    analyzer.set_z3_timeout_ms(1000)
+
+    expr = implies(tirx.all(a > 0, b > 0, c > 0), ((b - a) // c) * c + a <= b)
+    solver = z3.Solver()
+    solver.from_string(analyzer.get_smtlib2(expr))
+    assert solver.check() == z3.unsat
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/ci/__init__.py

@@ -0,0 +1,17 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Infrastructure and tests for CI scripts"""

tests/python/ci/sample_prs/pr10786.json

@@ -0,0 +1,130 @@
+{
+  "title": "[Hexagon] 2-d allocation cleanup",
+  "body":
+      "- Added device validity check in allocation. HexagonDeviceAPI should only be called for CPU/Hexagon types.\r\n\r\n- Check for \"global.vtcm\" scope instead of \"vtcm\".  The ccope of N-d allocations produced by `LowerVtcmAlloc` should be `\"global.vtcm\"`.  The previous check allowed unsupported scope such as `\"local.vtcm\"`.\r\n\r\n- Remove `vtcmallocs` entry after calling free.\n\n\nThanks for contributing to TVM!   Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from [Reviewers](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @ them in the pull request thread.\n\n\nPreviously, the vtcm allocation map kept dangling pointers to `HexagonBuffer` objects after they had been freed.\r\n\r\n- Rename N-d alloc and free packed functions.  Since most of the similar device functions use snake case, renaming `*.AllocND` to `*.alloc_nd` and `*.FreeND` to `*.free_nd`.\n\n\ncc @someone\n\r\n\r\nCo-authored-by: Adam Straw <astraw@octoml.ai>\n\n\nThanks for contributing to TVM!   Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from [Reviewers](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @ them in the pull request thread.\n\n",
+  "state": "OPEN",
+  "author": {
+    "login": "abc"
+  },
+  "comments": {
+    "pageInfo": {
+      "hasPreviousPage": false
+    },
+    "nodes": []
+  },
+  "authorCommits": {
+    "nodes": [
+      {
+        "commit": {
+          "authors": {
+            "nodes": [
+              {
+                "name": "Eric Lunderberg",
+                "email": "elunderberg@octoml.ai"
+              },
+              {
+                "name": "Adam Straw",
+                "email": "astraw@octoml.ai"
+              }
+            ]
+          }
+        }
+      }
+    ]
+  },
+  "commits": {
+    "nodes": [
+      {
+        "commit": {
+          "oid": "6f04bcf57d07f915a98fd91178f04d9e92a09fcd",
+          "statusCheckRollup": {
+            "contexts": {
+              "pageInfo": {
+                "hasNextPage": false
+              },
+              "nodes": [
+                {
+                  "name": "MacOS",
+                  "checkSuite": {
+                    "workflowRun": {
+                      "workflow": {
+                        "name": "CI"
+                      }
+                    }
+                  },
+                  "status": "COMPLETED",
+                  "conclusion": "SUCCESS",
+                  "url": "https://github.com/apache/tvm/runs/5694945392"
+                },
+                {
+                  "name": "cc-reviewers",
+                  "checkSuite": {
+                    "workflowRun": {
+                      "workflow": {
+                        "name": "PR"
+                      }
+                    }
+                  },
+                  "status": "COMPLETED",
+                  "conclusion": "SUCCESS",
+                  "url": "https://github.com/apache/tvm/runs/5694945029"
+                },
+                {
+                  "name": "tag-teams",
+                  "checkSuite": {
+                    "workflowRun": {
+                      "workflow": {
+                        "name": "Teams"
+                      }
+                    }
+                  },
+                  "status": "COMPLETED",
+                  "conclusion": "SUCCESS",
+                  "url": "https://github.com/apache/tvm/runs/5694945030"
+                },
+                {
+                  "name": "Windows",
+                  "checkSuite": {
+                    "workflowRun": {
+                      "workflow": {
+                        "name": "CI"
+                      }
+                    }
+                  },
+                  "status": "COMPLETED",
+                  "conclusion": "SUCCESS",
+                  "url": "https://github.com/apache/tvm/runs/5694945524"
+                },
+                {
+                  "state": "SUCCESS",
+                  "context": "tvm-ci/pr-head",
+                  "targetUrl": "https://ci.tlcpack.ai/job/tvm/job/PR-10786/1/display/redirect"
+                }
+              ]
+            }
+          }
+        }
+      }
+    ]
+  },
+  "reviewDecision": "APPROVED",
+  "reviews": {
+    "pageInfo": {
+      "hasPreviousPage": false
+    },
+    "nodes": [
+      {
+        "body": "@tvm-bot merge",
+        "updatedAt": "2022-03-25T22:13:50Z",
+        "authorCanPushToRepository": true,
+        "commit": {
+          "oid": "6f04bcf57d07f915a98fd91178f04d9e92a09fcd"
+        },
+        "author": {
+          "login": "kparzysz-quic"
+        },
+        "state": "APPROVED"
+      }
+    ]
+  }
+}

tests/python/ci/test_ci.py

@@ -0,0 +1,422 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: RUF005
+"""Test various CI scripts and GitHub Actions workflows"""
+
+import json
+import subprocess
+from pathlib import Path
+
+import pytest
+
+import tvm.testing
+
+from .test_utils import JENKINS_SCRIPT_ROOT, TempGit, run_script
+
+
+def parameterize_named(**kwargs):
+    keys = next(iter(kwargs.values())).keys()
+    return pytest.mark.parametrize(
+        ",".join(keys), [tuple(d.values()) for d in kwargs.values()], ids=kwargs.keys()
+    )
+
+
+@parameterize_named(
+    dont_skip_main=dict(
+        commands=[],
+        should_skip=False,
+        pr_title="[skip ci] test",
+        why="ci should not be skipped on main",
+    ),
+    dont_skip_main_with_commit=dict(
+        commands=[
+            ["commit", "--allow-empty", "--message", "[skip ci] commit 1"],
+        ],
+        should_skip=False,
+        pr_title="[skip ci] test",
+        why="ci should not be skipped on main",
+    ),
+    skip_on_new_branch=dict(
+        commands=[
+            ["checkout", "-b", "some_new_branch"],
+            ["commit", "--allow-empty", "--message", "[skip ci] commit 1"],
+        ],
+        should_skip=True,
+        pr_title="[skip ci] test",
+        why="ci should be skipped on a branch with [skip ci] in the last commit",
+    ),
+    no_skip_in_pr_title=dict(
+        commands=[
+            ["checkout", "-b", "some_new_branch"],
+            ["commit", "--allow-empty", "--message", "[skip ci] commit 1"],
+        ],
+        should_skip=False,
+        pr_title="[no skip ci] test",
+        why="ci should not be skipped on a branch with "
+        "[skip ci] in the last commit but not the PR title",
+    ),
+    skip_in_pr_title=dict(
+        commands=[
+            ["checkout", "-b", "some_new_branch"],
+            ["commit", "--allow-empty", "--message", "[skip ci] commit 1"],
+            ["commit", "--allow-empty", "--message", "commit 2"],
+        ],
+        should_skip=True,
+        pr_title="[skip ci] test",
+        why="ci should be skipped with [skip ci] in the PR title",
+    ),
+    skip_in_pr_title_many_commits=dict(
+        commands=[
+            ["checkout", "-b", "some_new_branch"],
+            ["commit", "--allow-empty", "--message", "commit 1"],
+            ["commit", "--allow-empty", "--message", "commit 2"],
+            ["commit", "--allow-empty", "--message", "commit 3"],
+            ["commit", "--allow-empty", "--message", "commit 4"],
+        ],
+        should_skip=True,
+        pr_title="[skip ci] test",
+        why="ci should be skipped with [skip ci] in the PR title",
+    ),
+    skip_anywhere_in_title=dict(
+        commands=[
+            ["checkout", "-b", "some_new_branch"],
+        ],
+        should_skip=True,
+        pr_title="[something][skip ci] test",
+        why="skip ci tag should work anywhere in title",
+    ),
+)
+def test_skip_ci(tmpdir_factory, commands, should_skip, pr_title, why):
+    """
+    Test that CI is skipped when it should be
+    """
+    skip_ci_script = JENKINS_SCRIPT_ROOT / "git_skip_ci.py"
+
+    git = TempGit(tmpdir_factory.mktemp("tmp_git_dir"))
+
+    git.run("config", "user.name", "ci")
+    git.run("config", "user.email", "email@example.com")
+    git.run("commit", "--allow-empty", "--message", "base commit")
+    for command in commands:
+        git.run(*command)
+    pr_number = "1234"
+    proc = run_script(
+        [skip_ci_script, "--pr", pr_number, "--pr-title", pr_title],
+        cwd=git.cwd,
+        check=False,
+    )
+    expected = 0 if should_skip else 1
+    if proc.returncode != expected:
+        raise RuntimeError(
+            f"Unexpected return code {proc.returncode} "
+            f"(expected {expected}) in {why}:\n{proc.stdout}"
+        )
+
+
+@parameterize_named(
+    no_file=dict(files=[], should_skip=True),
+    readme=dict(files=["README.md"], should_skip=True),
+    c_file=dict(files=["test.c"], should_skip=False),
+    c_and_readme=dict(files=["test.c", "README.md"], should_skip=False),
+    src_file_and_readme=dict(
+        files=["src/autotvm/feature_visitor.cc", "README.md"], should_skip=False
+    ),
+    yaml_and_readme=dict(files=[".asf.yaml", "docs/README.md"], should_skip=True),
+)
+def test_skip_globs(tmpdir_factory, files, should_skip):
+    """
+    Test that CI is skipped if only certain files are edited
+    """
+    script = JENKINS_SCRIPT_ROOT / "git_skip_ci_globs.py"
+
+    git = TempGit(tmpdir_factory.mktemp("tmp_git_dir"))
+
+    proc = run_script(
+        [
+            script,
+            "--files",
+            ",".join(files),
+        ],
+        check=False,
+        cwd=git.cwd,
+    )
+
+    if should_skip:
+        assert proc.returncode == 0
+    else:
+        assert proc.returncode == 1
+
+
+def assert_in(needle: str, haystack: str):
+    """
+    Check that 'needle' is in 'haystack'
+    """
+    if needle not in haystack:
+        raise AssertionError(f"item not found:\n{needle}\nin:\n{haystack}")
+
+
+@tvm.testing.skip_if_wheel_test
+@parameterize_named(
+    same_tags=dict(
+        tlcpackstaging_body={
+            "results": [
+                {
+                    "last_updated": "2022-06-01T00:00:00.123456Z",
+                    "name": "123-123-abc",
+                },
+            ]
+        },
+        tlcpack_body={
+            "results": [
+                {
+                    "last_updated": "2022-06-01T00:00:00.123456Z",
+                    "name": "123-123-abc",
+                },
+            ]
+        },
+        expected="Tag names were the same, no update needed",
+        expected_images=[],
+    ),
+    staging_update=dict(
+        tlcpackstaging_body={
+            "results": [
+                {
+                    "last_updated": "2022-06-01T01:00:00.123456Z",
+                    "name": "234-234-abc-staging",
+                },
+                {
+                    "last_updated": "2022-06-01T00:00:00.123456Z",
+                    "name": "456-456-abc",
+                },
+            ]
+        },
+        tlcpack_body={
+            "results": [
+                {
+                    "last_updated": "2022-06-01T00:00:00.123456Z",
+                    "name": "123-123-abc",
+                },
+            ]
+        },
+        expected="Using tlcpackstaging tag on tlcpack",
+        expected_images=[
+            "ci_arm: tlcpack/ci-arm:456-456-abc",
+        ],
+    ),
+    tlcpack_update=dict(
+        tlcpackstaging_body={
+            "results": [
+                {
+                    "last_updated": "2022-06-01T00:00:00.123456Z",
+                    "name": "123-123-abc",
+                },
+            ]
+        },
+        tlcpack_body={
+            "results": [
+                {
+                    "last_updated": "2022-06-01T00:01:00.123456Z",
+                    "name": "234-234-abc",
+                },
+            ]
+        },
+        expected="Found newer image, using: tlcpack",
+        expected_images=[
+            "ci_arm: tlcpack/ci-arm:234-234-abc",
+        ],
+    ),
+)
+def test_open_docker_update_pr(
+    tmpdir_factory, tlcpackstaging_body, tlcpack_body, expected, expected_images
+):
+    """Test workflow to open a PR to update Docker images"""
+    tag_script = JENKINS_SCRIPT_ROOT / "open_docker_update_pr.py"
+
+    git = TempGit(tmpdir_factory.mktemp("tmp_git_dir"))
+    git.run("config", "user.name", "ci")
+    git.run("config", "user.email", "email@example.com")
+    images = [
+        "ci_arm",
+        "ci_cortexm",
+        "ci_cpu",
+        "ci_gpu",
+        "ci_minimal",
+        "ci_riscv",
+        "ci_wasm",
+    ]
+
+    docker_data = {}
+    for image in images:
+        docker_data[f"repositories/tlcpackstaging/{image}/tags"] = tlcpackstaging_body
+        docker_data[f"repositories/tlcpack/{image.replace('_', '-')}/tags"] = tlcpack_body
+
+    proc = run_script(
+        [
+            tag_script,
+            "--dry-run",
+            "--testing-docker-data",
+            json.dumps(docker_data),
+        ],
+        cwd=git.cwd,
+        env={"GITHUB_TOKEN": "1234"},
+        stderr=subprocess.STDOUT,
+    )
+
+    for line in expected_images:
+        if line not in proc.stdout:
+            raise RuntimeError(f"Missing line {line} in output:\n{proc.stdout}")
+
+    assert_in(expected, proc.stdout)
+
+
+@parameterize_named(
+    use_tlcpack=dict(
+        images=["ci_arm", "ci_cpu"],
+        expected={
+            "ci_arm": "tlcpack/ci-arm:abc-abc-123",
+            "ci_cpu": "tlcpack/ci-cpu:abc-abc-234",
+        },
+    ),
+    use_staging=dict(
+        images=["ci_arm2"],
+        expected={
+            "ci_arm2": "tlcpackstaging/ci_arm2:abc-abc-123",
+        },
+    ),
+)
+def test_determine_docker_images(tmpdir_factory, images, expected):
+    """Test script to decide whether to use tlcpack or tlcpackstaging for images"""
+    script = JENKINS_SCRIPT_ROOT / "determine_docker_images.py"
+
+    git_dir = tmpdir_factory.mktemp("tmp_git_dir")
+
+    docker_data = {
+        "repositories/tlcpack/ci-arm/tags/abc-abc-123": {},
+        "repositories/tlcpack/ci-cpu/tags/abc-abc-234": {},
+    }
+
+    images_data = {
+        "ci_arm": "tlcpack/ci-arm:abc-abc-123",
+        "ci_cpu": "tlcpack/ci-cpu:abc-abc-234",
+        "ci_arm2": "tlcpack/ci-arm2:abc-abc-123",
+    }
+
+    run_script(
+        [
+            script,
+            "--testing-docker-data",
+            json.dumps(docker_data),
+            "--testing-images-data",
+            json.dumps(images_data),
+            "--base-dir",
+            git_dir,
+        ]
+        + images,
+        cwd=git_dir,
+    )
+
+    for expected_filename, expected_image in expected.items():
+        with open(Path(git_dir) / expected_filename) as f:
+            actual_image = f.read()
+
+        assert actual_image == expected_image
+
+
+@parameterize_named(
+    invalid_name=dict(
+        changed_files=[],
+        name="abc",
+        check="Image abc is not using new naming scheme",
+        expected_code=1,
+    ),
+    no_hash=dict(
+        changed_files=[], name="123-123-abc", check="No extant hash found", expected_code=1
+    ),
+    no_changes=dict(
+        changed_files=[["test.txt"]],
+        name=None,
+        check="Did not find changes, no rebuild necessary",
+        expected_code=0,
+    ),
+    docker_changes=dict(
+        changed_files=[["test.txt"], ["docker/test.txt"]],
+        name=None,
+        check="Found docker changes",
+        expected_code=2,
+    ),
+)
+def test_should_rebuild_docker(tmpdir_factory, changed_files, name, check, expected_code):
+    """
+    Check that the Docker images are built when necessary
+    """
+    tag_script = JENKINS_SCRIPT_ROOT / "should_rebuild_docker.py"
+
+    git = TempGit(tmpdir_factory.mktemp("tmp_git_dir"))
+    git.run("config", "user.name", "ci")
+    git.run("config", "user.email", "email@example.com")
+
+    git_path = Path(git.cwd)
+    for i, commits in enumerate(changed_files):
+        for filename in commits:
+            path = git_path / filename
+            path.parent.mkdir(exist_ok=True, parents=True)
+            path.touch()
+            git.run("add", filename)
+
+        git.run("commit", "-m", f"message {i}")
+
+    if name is None:
+        ref = "HEAD"
+        if len(changed_files) > 1:
+            ref = f"HEAD~{len(changed_files) - 1}"
+        proc = git.run("rev-parse", ref, stdout=subprocess.PIPE)
+        last_hash = proc.stdout.strip()
+        name = f"123-123-{last_hash}"
+
+    docker_data = {
+        "repositories/tlcpack": {
+            "results": [
+                {
+                    "name": "ci-something",
+                },
+                {
+                    "name": "something-else",
+                },
+            ],
+        },
+        "repositories/tlcpack/ci-something/tags": {
+            "results": [{"name": name}, {"name": name + "old"}],
+        },
+    }
+
+    proc = run_script(
+        [
+            tag_script,
+            "--testing-docker-data",
+            json.dumps(docker_data),
+        ],
+        stderr=subprocess.STDOUT,
+        cwd=git.cwd,
+        check=False,
+    )
+
+    assert_in(check, proc.stdout)
+    assert proc.returncode == expected_code
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/ci/test_tvmbot.py

@@ -0,0 +1,281 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501
+"""
+Test the @tvm-bot merge code
+"""
+
+import json
+from pathlib import Path
+from typing import Any
+
+import tvm
+
+from .test_utils import GITHUB_SCRIPT_ROOT, TempGit, run_script
+
+SUCCESS_EXPECTED_OUTPUT = """
+Dry run, would have merged with url=pulls/10786/merge and data={
+  "commit_title": "[Hexagon] 2-d allocation cleanup (#10786)",
+  "commit_message": "- Added device validity check in allocation. HexagonDeviceAPI should only be called for CPU/Hexagon types.\\n\\n- Check for \\"global.vtcm\\" scope instead of \\"vtcm\\".  The ccope of N-d allocations produced by `LowerVtcmAlloc` should be `\\"global.vtcm\\"`.  The previous check allowed unsupported scope such as `\\"local.vtcm\\"`.\\n\\n- Remove `vtcmallocs` entry after calling free.\\n\\nPreviously, the vtcm allocation map kept dangling pointers to `HexagonBuffer` objects after they had been freed.\\n\\n- Rename N-d alloc and free packed functions.  Since most of the similar device functions use snake case, renaming `*.AllocND` to `*.alloc_nd` and `*.FreeND` to `*.free_nd`.\\n\\n\\ncc someone\\n\\n\\nCo-authored-by: Adam Straw <astraw@octoml.ai>",
+  "sha": "6f04bcf57d07f915a98fd91178f04d9e92a09fcd",
+  "merge_method": "squash"
+}
+""".strip()
+
+
+class _TvmBotTest:
+    NUMBER = 10786
+
+    def preprocess_data(self, data: dict[str, Any]):
+        """
+        Used to pre-process PR data before running the test. Override as
+        necessary to edit data for specific test cases.
+        """
+        return data
+
+    @tvm.testing.skip_if_wheel_test
+    def test(self, tmpdir_factory):
+        """
+        Run the tvm-bot script using the data from preprocess_data
+        """
+        mergebot_script = GITHUB_SCRIPT_ROOT / "github_tvmbot.py"
+        test_json_dir = Path(__file__).resolve().parent / "sample_prs"
+        with open(test_json_dir / f"pr{self.NUMBER}.json") as f:
+            test_data = json.load(f)
+
+        # Update testing data with replacements / additions
+        test_data = self.preprocess_data(test_data)
+
+        git = TempGit(tmpdir_factory.mktemp("tmp_git_dir"))
+
+        comment = {
+            "body": self.COMMENT,
+            "id": 123,
+            "user": {
+                "login": self.USER,
+            },
+        }
+        allowed_users = [{"login": "abc"}, {"login": "other-abc"}]
+
+        proc = run_script(
+            [
+                mergebot_script,
+                "--pr",
+                self.NUMBER,
+                "--dry-run",
+                "--run-url",
+                "https://example.com",
+                "--testing-pr-json",
+                json.dumps(test_data),
+                "--testing-collaborators-json",
+                json.dumps(allowed_users),
+                "--testing-mentionable-users-json",
+                json.dumps(allowed_users),
+                "--trigger-comment-json",
+                json.dumps(comment),
+            ],
+            env={
+                "TVM_BOT_JENKINS_TOKEN": "123",
+                "GH_ACTIONS_TOKEN": "123",
+            },
+            cwd=git.cwd,
+        )
+
+        if self.EXPECTED not in proc.stderr:
+            raise RuntimeError(f"{proc.stderr}\ndid not contain\n{self.EXPECTED}")
+
+
+class TestNoRequest(_TvmBotTest):
+    """
+    A PR for which the mergebot runs but no merge is requested
+    """
+
+    COMMENT = "@tvm-bot do something else"
+    USER = "abc"
+    EXPECTED = "Command 'do something else' did not match anything"
+
+    def preprocess_data(self, data: dict[str, Any]):
+        data["reviews"]["nodes"][0]["body"] = "nothing"
+        return data
+
+
+class TestSuccessfulMerge(_TvmBotTest):
+    """
+    Everything is fine so this PR will merge
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "abc"
+    EXPECTED = SUCCESS_EXPECTED_OUTPUT
+
+
+class TestBadCI(_TvmBotTest):
+    """
+    A PR which failed CI and cannot merge
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "abc"
+    EXPECTED = "Cannot merge, these CI jobs are not successful on"
+
+    def preprocess_data(self, data: dict[str, Any]):
+        # Mark the Jenkins build as failed
+        contexts = data["commits"]["nodes"][0]["commit"]["statusCheckRollup"]["contexts"]["nodes"]
+        for context in contexts:
+            if "context" in context and context["context"] == "tvm-ci/pr-head":
+                context["state"] = "FAILED"
+        return data
+
+
+class TestOldReview(_TvmBotTest):
+    """
+    A PR with passing CI and approving reviews on an old commit so it cannot merge
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "abc"
+    EXPECTED = "Cannot merge, did not find any approving reviews"
+
+    def preprocess_data(self, data: dict[str, Any]) -> dict[str, Any]:
+        data["reviews"]["nodes"][0]["commit"]["oid"] = "abc12345"
+        return data
+
+
+class TestMissingJob(_TvmBotTest):
+    """
+    PR missing an expected CI job and cannot merge
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "abc"
+    EXPECTED = "Cannot merge, missing expected jobs"
+
+    def preprocess_data(self, data: dict[str, Any]) -> dict[str, Any]:
+        contexts = data["commits"]["nodes"][0]["commit"]["statusCheckRollup"]["contexts"]["nodes"]
+        for context in contexts:
+            if "context" in context and context["context"] == "tvm-ci/pr-head":
+                context["context"] = "something"
+        return data
+
+
+class TestInvalidAuthor(_TvmBotTest):
+    """
+    Merge requester is not a committer and cannot merge
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "not-abc"
+    EXPECTED = "Failed auth check 'collaborators', quitting"
+
+
+class TestUnauthorizedComment(_TvmBotTest):
+    """
+    Check that a merge comment not from a CONTRIBUTOR is rejected
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "not-abc2"
+    EXPECTED = "Failed auth check 'collaborators'"
+
+
+class TestNoReview(_TvmBotTest):
+    """
+    Check that a merge request without any reviews is rejected
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "abc"
+    EXPECTED = "Cannot merge, did not find any approving reviews from users with write access"
+
+    def preprocess_data(self, data: dict[str, Any]) -> dict[str, Any]:
+        data["reviews"]["nodes"] = []
+        return data
+
+
+class TestChangesRequested(_TvmBotTest):
+    """
+    Check that a merge request with a 'Changes Requested' review is rejected
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "abc"
+    EXPECTED = "Cannot merge, found [this review]"
+
+    def preprocess_data(self, data: dict[str, Any]) -> dict[str, Any]:
+        data["reviews"]["nodes"][0]["state"] = "CHANGES_REQUESTED"
+        data["reviews"]["nodes"][0]["url"] = "http://example.com"
+        return data
+
+
+class TestCoAuthors(_TvmBotTest):
+    """
+    Check that a merge request with co-authors generates the correct commit message
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "abc"
+    EXPECTED = "Co-authored-by: Some One <someone@email.com>"
+
+    def preprocess_data(self, data: dict[str, Any]) -> dict[str, Any]:
+        data["authorCommits"]["nodes"][0]["commit"]["authors"]["nodes"].append(
+            {"name": "Some One", "email": "someone@email.com"}
+        )
+        return data
+
+
+class TestRerunCI(_TvmBotTest):
+    """
+    Start a new CI job
+    """
+
+    COMMENT = "@tvm-bot rerun"
+    USER = "abc"
+    EXPECTED = "Rerunning ci with"
+
+
+class TestRerunPermissions(_TvmBotTest):
+    """
+    Start a new CI job as an unauthorized user
+    """
+
+    COMMENT = "@tvm-bot rerun"
+    USER = "someone"
+    EXPECTED = "Failed auth check 'mentionable_users', quitting"
+
+
+class TestRerunNonAuthor(_TvmBotTest):
+    """
+    Start a new CI job as a mentionable user
+    """
+
+    COMMENT = "@tvm-bot rerun"
+    USER = "other-abc"
+    EXPECTED = "Passed auth check 'mentionable_users', continuing"
+
+
+class TestIgnoreJobs(_TvmBotTest):
+    """
+    Ignore GitHub Actions jobs that don't start with CI /
+    """
+
+    COMMENT = "@tvm-bot merge"
+    USER = "abc"
+    EXPECTED = "Dry run, would have merged"
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/ci/test_utils.py

@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: RUF005
+"""
+Constants used in various CI tests
+"""
+
+import os
+import pathlib
+import subprocess
+from typing import Any
+
+REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent.parent.parent
+GITHUB_SCRIPT_ROOT = REPO_ROOT / "ci" / "scripts" / "github"
+JENKINS_SCRIPT_ROOT = REPO_ROOT / "ci" / "scripts" / "jenkins"
+
+
+class TempGit:
+    """
+    A wrapper to run commands in a directory (specifically for use in CI tests)
+    """
+
+    def __init__(self, cwd):
+        self.cwd = cwd
+        # Jenkins git is too old and doesn't have 'git init --initial-branch',
+        # so init and checkout need to be separate steps
+        self.run("init", stderr=subprocess.PIPE, stdout=subprocess.PIPE)
+        self.run("checkout", "-b", "main", stderr=subprocess.PIPE)
+        self.run("remote", "add", "origin", "https://github.com/apache/tvm.git")
+
+    def run(self, *args, **kwargs):
+        """
+        Run a git command based on *args
+        """
+        proc = subprocess.run(
+            ["git"] + list(args), encoding="utf-8", cwd=self.cwd, check=False, **kwargs
+        )
+        if proc.returncode != 0:
+            raise RuntimeError(f"git command failed: '{args}'")
+
+        return proc
+
+
+def run_script(command: list[Any], check: bool = True, **kwargs):
+    """
+    Wrapper to run a script and print its output if there was an error
+    """
+    command = [str(c) for c in command]
+    kwargs_to_send = {
+        "stdout": subprocess.PIPE,
+        "stderr": subprocess.PIPE,
+        "encoding": "utf-8",
+    }
+    env = kwargs.pop("env", None)
+    if env is not None:
+        kwargs_to_send["env"] = {**os.environ, **env}
+    kwargs_to_send.update(kwargs)
+    proc = subprocess.run(
+        command,
+        check=False,
+        **kwargs_to_send,
+    )
+    if check and proc.returncode != 0:
+        raise RuntimeError(f"Process failed:\nstdout:\n{proc.stdout}\n\nstderr:\n{proc.stderr}")
+
+    return proc

tests/python/codegen/test_codegen_assert.py

@@ -0,0 +1,124 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""AssertStmt codegen tests: verify kind and message_parts produce correct exceptions."""
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+
+codegen_target = tvm.testing.parameter("llvm", "c")
+
+
+def test_assert_runtime_error(codegen_target):
+    """AssertStmt with RuntimeError kind produces RuntimeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("RuntimeError", ["Expected non-null input"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(RuntimeError, match="Expected non-null input"):
+        lib(0)
+
+
+def test_assert_value_error(codegen_target):
+    """AssertStmt with ValueError kind produces ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("ValueError", ["Shape mismatch: expected 4 got 8"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(ValueError, match="Shape mismatch"):
+        lib(0)
+
+
+def test_assert_type_error(codegen_target):
+    """AssertStmt with TypeError kind produces TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("TypeError", ["Expected Tensor but got int"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(TypeError, match="Expected Tensor but got int"):
+        lib(0)
+
+
+def test_assert_multi_part_message(codegen_target):
+    """Multi-part messages are correctly concatenated at runtime."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("ValueError", ["Expected shape ", "4", " but got ", "8"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(ValueError, match="Expected shape 4 but got 8"):
+        lib(0)
+
+
+def test_assert_passing_condition(codegen_target):
+    """Passing assertion does not raise."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("RuntimeError", ["This should not be raised"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    lib(1)  # should pass without error
+
+
+def test_assert_many_parts(codegen_target):
+    """Assertion with 8 parts concatenated correctly."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("RuntimeError", ["p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(RuntimeError, match="p0p1p2p3p4p5p6p7"):
+        lib(0)
+
+
+def test_tvmscript_assert_preserves_kind(codegen_target):
+    """Regression: TVMScript structured assert preserves kind at runtime."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("ValueError", ["x must be positive"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(ValueError, match="x must be positive"):
+        lib(0)
+
+
+def test_tvmscript_assert_preserves_parts(codegen_target):
+    """Regression: TVMScript structured assert with separate parts."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("ValueError", ["x must be ", "positive"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(ValueError, match="x must be positive"):
+        lib(0)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_codegen_error_handling.py

@@ -0,0 +1,471 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Runtime error message tests for MakePackedAPI + TVMFFIABIBuilder.
+
+All tests compile TVMScript functions and verify the correct Python exception
+type and exact error message at runtime.
+"""
+
+import re
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+# Parameterize over both LLVM and C backends
+codegen_target = tvm.testing.parameter("llvm", "c")
+
+
+# ── Argument count errors ────────────────────────────────────
+
+
+def test_wrong_argument_count_error(codegen_target):
+    """Wrong argument count produces TypeError with function signature."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        n0 = T.int64()
+        A = T.match_buffer(a, (n0,), "float32")
+        B = T.match_buffer(b, (n0,), "float32")
+        for i in range(n0):
+            B[i] = A[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a, b)  # correct input should pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Expected 2 arguments when calling:\n"
+            "  `func(A: Tensor([n0], float32), B: Tensor([n0], float32))`"
+        ),
+    ):
+        lib()
+
+
+# ── Type mismatch errors (tensor parameters) ────────────────
+
+
+def test_type_mismatch_non_tensor(codegen_target):
+    """Passing a non-tensor where a tensor is expected raises TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        n0 = T.int64()
+        A = T.match_buffer(a, (n0,), "float32")
+        B = T.match_buffer(b, (n0,), "float32")
+        for i in range(n0):
+            B[i] = A[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a, b)  # correct input should pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #1 when calling:\n"
+            "  `func(A: Tensor([n0], float32), B: Tensor([n0], float32))`,\n"
+            "  expected Tensor"
+        ),
+    ):
+        lib(a, 1)
+
+
+# ── Shape mismatch errors ───────────────────────────────────
+
+
+def test_shape_mismatch_shared_variable(codegen_target):
+    """b has different shape than a when they share symbolic variable n0."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        n0 = T.int64()
+        A = T.match_buffer(a, (n0,), "float32")
+        B = T.match_buffer(b, (n0,), "float32")
+        for i in range(n0):
+            B[i] = A[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a, b)  # correct input should pass
+
+    b_short = tvm.runtime.tensor(np.zeros(126, dtype="float32"))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Mismatched B.shape[0] on argument #1 when calling:\n"
+            "  `func(A: Tensor([n0], float32), B: Tensor([n0], float32))`,\n"
+            "  expected to match A.shape[0]"
+        ),
+    ):
+        lib(a, b_short)
+
+
+def test_invalid_shape_fixed(codegen_target):
+    """Passing wrong shape for a fixed buffer dimension raises ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((128,), "float32"), b: T.Buffer((128,), "float32")):
+        for i in range(128):
+            b[i] = a[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a, b)  # correct input should pass
+
+    a_wrong = tvm.runtime.tensor(np.zeros(256, dtype="float32"))
+    b_wrong = tvm.runtime.tensor(np.zeros(256, dtype="float32"))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Invalid a.shape[0] on argument #0 when calling:\n"
+            "  `func(a: Tensor([128], float32), b: Tensor([128], float32))`,\n"
+            "  expected 128"
+        ),
+    ):
+        lib(a_wrong, b_wrong)
+
+
+# ── ndim mismatch errors ────────────────────────────────────
+
+
+def test_ndim_mismatch_error(codegen_target):
+    """ndim mismatch produces ValueError with function signature."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((4, 8), "float32"), b: T.Buffer((4, 8), "float32")):
+        for i, j in T.grid(4, 8):
+            b[i, j] = a[i, j]
+
+    lib = tvm.compile(func, target=codegen_target)
+    a_ok = tvm.runtime.tensor(np.zeros((4, 8), dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros((4, 8), dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(4, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(4, dtype="float32"))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Mismatched a.ndim on argument #0 when calling:\n"
+            "  `func(a: Tensor([4, 8], float32), b: Tensor([4, 8], float32))`,\n"
+            "  expected 2"
+        ),
+    ):
+        lib(a, b)
+
+
+# ── dtype mismatch errors ───────────────────────────────────
+
+
+def test_dtype_mismatch_error(codegen_target):
+    """dtype mismatch produces TypeError with function signature."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((8,), "float32"), b: T.Buffer((8,), "float32")):
+        for i in range(8):
+            b[i] = a[i]
+
+    lib = tvm.compile(func, target=codegen_target)
+    a_ok = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(8, dtype="int32"))
+    b = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched a.dtype on argument #0 when calling:\n"
+            "  `func(a: Tensor([8], float32), b: Tensor([8], float32))`,\n"
+            "  expected float32"
+        ),
+    ):
+        lib(a, b)
+
+
+# ── Data alignment errors ──────────────────────────────────
+
+
+@pytest.mark.skip(reason="alignment check disabled for now, revisit after merge")
+def test_data_alignment_error(codegen_target):
+    """Misaligned buffer data pointer raises ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((128,), "float32"), b: T.Buffer((128,), "float32")):
+        for i in range(128):
+            b[i] = a[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    # Slice off first element of a 129-element array to create misaligned data pointer
+    np_arr = np.zeros(129, dtype="float32")
+    a_misaligned = tvm_ffi.from_dlpack(np_arr[1:])
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Misaligned Tensor data on argument #0 when calling:\n"
+            "  `func(a: Tensor([128], float32), b: Tensor([128], float32))`,\n"
+            "  expected data alignment=64 bytes"
+        ),
+    ):
+        lib(a_misaligned, b)
+
+
+# ── Compact strides mismatch errors ────────────────────────
+
+
+def test_strides_mismatch_transposed(codegen_target):
+    """Transposed (non-compact) strides raise ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((128, 128), "float32"), b: T.Buffer((128, 128), "float32")):
+        for i, j in T.grid(128, 128):
+            b[i, j] = a[i, j] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a_ok = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    # Use Fortran-order array to get non-compact (non-C-contiguous) strides
+    np_arr = np.asfortranarray(np.zeros((128, 128), dtype="float32"))
+    a_transposed = tvm_ffi.from_dlpack(np_arr)
+    b = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Mismatched a.strides on argument #0 when calling:\n"
+            "  `func(a: Tensor([128, 128], float32), b: Tensor([128, 128], float32))`,\n"
+            "  expected to be compact array"
+        ),
+    ):
+        lib(a_transposed, b)
+
+
+# ── Device mismatch errors ─────────────────────────────────
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_device_mismatch_error():
+    """Passing GPU tensor to CPU function raises ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((128,), "float32"), b: T.Buffer((128,), "float32")):
+        for i in range(128):
+            b[i] = a[i] + T.float32(1)
+
+    lib = tvm.compile(func, target="llvm")
+    a_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    def run_and_check():
+        a_gpu = tvm.runtime.tensor(np.zeros(128, dtype="float32"), device=tvm.cuda(0))
+        b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "Mismatched a.device_type on argument #0 when calling:\n"
+                "  `func(a: Tensor([128], float32), b: Tensor([128], float32))`,\n"
+                "  expected cpu"
+            ),
+        ):
+            lib(a_gpu, b)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ── Scalar type mismatch errors ─────────────────────────────
+
+
+def test_type_mismatch_int_parameter(codegen_target):
+    """Passing a tensor where an int is expected raises TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32) -> T.int32:
+        if x > 0:
+            return 10
+        else:
+            return 20
+
+    lib = tvm.compile(func, target=codegen_target)
+    assert lib(5) == 10  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #0 when calling:\n  `func(x: int32)`,\n  expected int"
+        ),
+    ):
+        lib(a)
+
+
+def test_type_mismatch_float_parameter(codegen_target):
+    """Passing a tensor where a float is expected raises TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.float32) -> T.int32:
+        if x > T.float32(0):
+            return 1
+        else:
+            return 0
+
+    lib = tvm.compile(func, target=codegen_target)
+    assert lib(1.0) == 1  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #0 when calling:\n  `func(x: float32)`,\n  expected float"
+        ),
+    ):
+        lib(a)
+
+
+def test_type_mismatch_bool_parameter(codegen_target):
+    """Passing a tensor where a bool is expected raises TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.bool) -> T.int32:
+        if x:
+            return 1
+        else:
+            return 0
+
+    lib = tvm.compile(func, target=codegen_target)
+    assert lib(True) == 1  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #0 when calling:\n  `func(x: bool)`,\n  expected boolean"
+        ),
+    ):
+        lib(a)
+
+
+# ── Forward-reference symbolic variable ────────────────────
+
+
+def test_forward_reference_symbolic_shape(codegen_target):
+    """Buffers sharing a symbolic var with forward reference compile and run correctly.
+
+    When buffer A has shape (batch_size+1,) and buffer B has shape (batch_size,),
+    batch_size is referenced in A's shape assertion before it is defined from B.
+    The three-sequence separation ensures this works. Also verifies the error
+    message uses rendered access paths (e.g. "B.shape[0] + 1") for shape checks.
+    """
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        batch_size = T.int64()
+        A = T.match_buffer(a, (batch_size + 1,), "int32")
+        B = T.match_buffer(b, (batch_size,), "int32")
+        for i in range(batch_size):
+            B[i] = A[i] + A[i + 1]
+
+    lib = tvm.compile(func, target=codegen_target)
+    # Correct inputs: A has shape (5,), B has shape (4,)
+    a = tvm.runtime.tensor(np.array([1, 2, 3, 4, 5], dtype="int32"))
+    b = tvm.runtime.tensor(np.zeros(4, dtype="int32"))
+    lib(a, b)
+    np.testing.assert_array_equal(b.numpy(), [3, 5, 7, 9])
+
+    # Wrong shape: A has shape (10,) but B has shape (4,), so batch_size=4 but A needs 5
+    a_wrong = tvm.runtime.tensor(np.zeros(10, dtype="int32"))
+    b_ok = tvm.runtime.tensor(np.zeros(4, dtype="int32"))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Invalid A.shape[0] on argument #0 when calling:\n"
+            "  `func(A: Tensor([batch_size + T.int64(1)], int32),"
+            " B: Tensor([batch_size], int32))`,\n"
+            "  expected B.shape[0] + 1"
+        ),
+    ):
+        lib(a_wrong, b_ok)
+
+
+# ── Mixed parameter type errors ────────────────────────────
+
+
+def test_invalid_arguments_mixed_params(codegen_target):
+    """Mixed bool + tensor function: type, dtype, and shape errors."""
+
+    @T.prim_func(s_tir=True)
+    def func(a0: T.bool, a1: T.Buffer([10], "float32")) -> T.int32:
+        return 0
+
+    lib = tvm.compile(func, target=codegen_target)
+    lib(True, tvm.runtime.tensor(np.zeros(10, dtype="float32")))  # correct input should pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #1 when calling:\n"
+            "  `func(a0: bool, a1: Tensor([10], float32))`,\n"
+            "  expected Tensor"
+        ),
+    ):
+        lib(1, 1)
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched a1.dtype on argument #1 when calling:\n"
+            "  `func(a0: bool, a1: Tensor([10], float32))`,\n"
+            "  expected float32"
+        ),
+    ):
+        lib(1, tvm.runtime.empty([10], "int32"))
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Invalid a1.shape[0] on argument #1 when calling:\n"
+            "  `func(a0: bool, a1: Tensor([10], float32))`,\n"
+            "  expected 10"
+        ),
+    ):
+        lib(False, tvm.runtime.empty([11], "float32"))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_gpu_codegen_allreduce.py

@@ -0,0 +1,174 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E741
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def _reduce_sum_module(d1, d2, d3):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1, d1, d2, d3), "float32"), B: T.Buffer((1, d1, d2), "float32")):
+            for i in T.thread_binding(1, thread="blockIdx.x"):
+                for j in T.thread_binding(d1, thread="threadIdx.z"):
+                    for k in T.thread_binding(d2, thread="threadIdx.y"):
+                        for l in T.thread_binding(d3, thread="threadIdx.x"):
+                            with T.sblock("reduce"):
+                                vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
+                                T.reads(A[vi, vj, vk, vl])
+                                T.writes(B[vi, vj, vk])
+                                with T.init():
+                                    B[vi, vj, vk] = T.float32(0.0)
+                                B[vi, vj, vk] = B[vi, vj, vk] + A[vi, vj, vk, vl]
+
+    return Module
+
+
+def _reduce_max_module(d1, d2, d3):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1, d1, d2, d3), "float32"), B: T.Buffer((1, d1, d2), "float32")):
+            for i in T.thread_binding(1, thread="blockIdx.x"):
+                for j in T.thread_binding(d1, thread="threadIdx.z"):
+                    for k in T.thread_binding(d2, thread="threadIdx.y"):
+                        for l in T.thread_binding(d3, thread="threadIdx.x"):
+                            with T.sblock("reduce"):
+                                vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
+                                T.reads(A[vi, vj, vk, vl])
+                                T.writes(B[vi, vj, vk])
+                                with T.init():
+                                    B[vi, vj, vk] = T.float32(-3.4028234663852886e38)
+                                B[vi, vj, vk] = T.max(B[vi, vj, vk], A[vi, vj, vk, vl])
+
+    return Module
+
+
+def generate_param_sets():
+    for d1 in range(1, 5):
+        for d2 in range(1, 5):
+            for d3 in [2, 4, 8, 12, 16, 32, 48, 64, 100, 128, 201, 256, 512, 1024]:
+                if d1 * d2 * d3 < 1024:
+                    yield (d1, d2, d3)
+
+
+dims = tvm.testing.parameter(*generate_param_sets())
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+    ],
+)
+def test_allreduce_sum(dims, target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    d1, d2, d3 = dims
+    mod = _reduce_sum_module(d1, d2, d3)
+    f = tvm.compile(mod, target=target)
+
+    # prepare input and output array
+    a_np = np.random.rand(1, d1, d2, d3).astype("float32")
+    b_np = a_np.sum(axis=-1).astype("float32")
+
+    def run_and_check():
+        dev = tvm.device(target)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(np.zeros_like(b_np), dev)
+        f(a, b)
+        tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-6, atol=1e-6)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+define_metal_compile_callback = tvm.testing.parameter(True, False)
+
+
+@pytest.fixture
+def optional_metal_compile_callback(define_metal_compile_callback):
+    name = "tvm_callback_metal_compile"
+    cached = tvm.get_global_func(name, allow_missing=True)
+
+    if define_metal_compile_callback:
+
+        @tvm.register_global_func(name, override=True)
+        def compile_metal(src, target):
+            from tvm.support.xcode import compile_metal  # pylint: disable=import-outside-toplevel
+
+            return compile_metal(src, sdk="macosx")
+
+    yield
+
+    if define_metal_compile_callback:
+        if cached is None:
+            tvm_ffi.registry.remove_global_func(name)
+        else:
+            tvm.register_global_func(name, cached, override=True)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_allreduce_sum_compile(optional_metal_compile_callback):
+    # Disable the parametrization over dims, at least for now
+    dims = (1, 1, 2)
+    target = "metal"
+
+    d1, d2, d3 = dims
+    mod = _reduce_sum_module(d1, d2, d3)
+    tvm.compile(mod, target=target)
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+    ],
+)
+def test_allreduce_max(dims, target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    d1, d2, d3 = dims
+    mod = _reduce_max_module(d1, d2, d3)
+    f = tvm.compile(mod, target=target)
+
+    # prepare input and output array
+    a_np = -np.random.rand(1, d1, d2, d3).astype("float32")
+    b_np = a_np.max(axis=-1).astype("float32")
+
+    def run_and_check():
+        dev = tvm.device(target)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(np.zeros_like(b_np), dev)
+        f(a, b)
+        tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-6, atol=1e-6)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_inject_ptx_ldg32.py

@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@T.prim_func(s_tir=True)
+def vector_add(A: T.Buffer((16), "float32"), B: T.Buffer((32), "float32")) -> None:
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    bx = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(bx, 1)
+    T.launch_thread(tx, 32)
+    with T.sblock():
+        A_local = T.sblock_alloc_buffer((32), "float32", scope="local")
+
+        with T.sblock():
+            T.reads(A[0:16])
+            T.writes(A_local[0:32])
+            A_local[tx] = T.if_then_else(tx % 2 == 0, A[tx // 2], T.float32(0), dtype="float32")
+            B[tx] = A_local[tx] + 1.0
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_inject_ptx_intrin():
+    f = vector_add
+    arch = tvm.support.nvcc.get_target_compute_version()
+    major, _ = tvm.support.nvcc.parse_compute_version(arch)
+    if major < 8:
+        # Require at least SM80
+        return
+    with tvm.transform.PassContext(config={"tirx.ptx.ldg32": True}):
+        mod = tvm.compile(f, target="cuda")
+    A_np = np.random.rand(16).astype("float32")
+    B_np = np.zeros(32).astype("float32")
+    C_np = np.zeros(32).astype("float32")
+
+    for i in range(32):
+        if i % 2 == 0:
+            C_np[i] = A_np[i // 2]
+        C_np[i] += 1.0
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_nd = tvm.runtime.tensor(A_np, device=dev)
+        B_nd = tvm.runtime.tensor(B_np, device=dev)
+        mod(A_nd, B_nd)
+        tvm.testing.assert_allclose(B_nd.numpy(), C_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    test_inject_ptx_intrin()

tests/python/codegen/test_target_codegen.py

@@ -0,0 +1,166 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+
+
+def test_buffer_store_predicate_not_supported():
+    target = "c"
+
+    @T.prim_func(s_tir=True)
+    def func(b: T.handle):
+        B = T.match_buffer(b, (8,), "float32")
+        B.vstore([T.Ramp(0, 2, 4)], T.Broadcast(1.0, 4), predicate=T.Broadcast(T.bool(True), 4))
+
+    err_msg = "Predicated buffer store is not supported."
+    with pytest.raises(RuntimeError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.compile(func)
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("opencl", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+        pytest.param("rocm", marks=pytest.mark.gpu),
+        pytest.param({"kind": "vulkan", "from_device": 0}, marks=pytest.mark.gpu),
+    ],
+)
+def test_buffer_store_predicate_not_supported_gpu(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (2, 3), "float32")
+        B = T.match_buffer(b, (6,), "float32")
+        T.func_attr({"global_symbol": "main"})
+        for i_0 in T.thread_binding(3, thread="threadIdx.x"):
+            B.vstore(
+                [T.Ramp(i_0, 1, 4)], T.Broadcast(1.0, 4), predicate=T.Broadcast(T.bool(True), 4)
+            )
+
+    err_msg = "Predicated buffer store is not supported."
+    with pytest.raises(RuntimeError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.compile(func)
+
+
+def test_buffer_load_predicate_not_supported():
+    target = "c"
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (8,), "float32")
+        B = T.match_buffer(b, (8,), "float32")
+        for i_0 in range(4):
+            B.vstore(
+                [T.Ramp(0, 2, 4)],
+                A.vload([T.Ramp(i_0, 1, 4)], predicate=T.Broadcast(T.bool(True), 4)),
+            )
+
+    err_msg = "Predicated buffer load is not supported."
+    with pytest.raises(RuntimeError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.compile(func)
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("opencl", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+        pytest.param("rocm", marks=pytest.mark.gpu),
+        pytest.param({"kind": "vulkan", "from_device": 0}, marks=pytest.mark.gpu),
+    ],
+)
+def test_buffer_load_predicate_not_supported_gpu(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (8,), "float32")
+        B = T.match_buffer(b, (8,), "float32")
+        for i_0 in T.thread_binding(3, thread="threadIdx.x"):
+            B.vstore(
+                [T.Ramp(0, 2, 4)],
+                A.vload([T.Ramp(i_0, 1, 4)], predicate=T.Broadcast(T.bool(True), 4)),
+            )
+
+    err_msg = "Predicated buffer load is not supported."
+    with pytest.raises(RuntimeError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.compile(func)
+
+
+@pytest.mark.parametrize("target", ["c", "llvm"])
+def test_codegen_loop_step(target):
+    if target != "c" and not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def test_loop_step(
+        A: T.Buffer((1024,), "float32"),
+        B: T.Buffer((1024,), "float32"),
+        C: T.Buffer((1024,), "float32"),
+    ):
+        for i in T.serial(3, 1024, step=96):
+            C[i] = A[i] + B[i]
+
+    with tvm.transform.PassContext(disabled_pass=["s_tir.CanonicalizeLoop"]):
+        lib = tvm.compile(test_loop_step, target=target)
+
+    src = lib.mod.inspect_source()
+    if target == "c":
+        assert src.find("for (int32_t i = 3; i < 1024; i += 96)") >= 0
+
+    dev = tvm.device(target, 0)
+    a_np = np.random.rand(1024).astype("float32")
+    b_np = np.random.rand(1024).astype("float32")
+    c_np = np.zeros(1024, dtype="float32")
+    a_tvm = tvm.runtime.tensor(a_np, dev)
+    b_tvm = tvm.runtime.tensor(b_np, dev)
+    c_tvm = tvm.runtime.tensor(c_np, dev)
+
+    lib(a_tvm, b_tvm, c_tvm)
+
+    c_result = c_tvm.numpy()
+
+    # Check that the loop executes at positions 3, 99, 195, 291, 387, 483, 579, 675, 771, 867, 963
+    for i in range(3, 1024, 96):
+        tvm.testing.assert_allclose(c_result[i], a_np[i] + b_np[i], rtol=1e-5)
+
+    # Assert non-touched positions remain zero
+    for i in range(0, 3):
+        assert c_result[i] == 0.0
+    for i in range(4, 1024):
+        if (i - 3) % 96 != 0:
+            assert c_result[i] == 0.0
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_aarch64.py

@@ -0,0 +1,647 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Codegen tests for AArch64
+"""
+
+import re
+
+import pytest
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.target.codegen import llvm_version_major
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_mul(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] * B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"mul\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_add(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] + B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"add\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_sub(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] - B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"sub\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_muladd(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle, var_D: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            D = T.match_buffer(var_D, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("D"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i], C[v_i])
+                    T.writes(D[v_i])
+                    D[v_i] = A[v_i] * B[v_i] + C[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        # Group the mad|mla alternation: a top-level alternation would let a bare
+        # "mad" match anywhere in the assembly (e.g. inside scalar "fmadd").
+        r"(?:mad|mla)\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]",
+        assembly,
+    )
+
+    if llvm_version_major() >= 18 and dtype in ("float", "float16"):
+        # Newer LLVM cost models (observed with LLVM 20) prefer a fixed-width
+        # NEON main loop over a scalable SVE loop for floating-point fmuladd
+        # on generic AArch64 targets, so also accept the NEON form.
+        loads += re.findall(r"ld[rp]\tq[0-9]", assembly)
+        matches += re.findall(r"fml[as]\tv[0-9]+\.[0-9]+[hs]", assembly)
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_max(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = T.max(A[v_i], B[v_i])
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    compare = re.findall(
+        r"cmgt\tp[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+    select = re.findall("sel\tz[0-9].[shdb], p[0-9], z[0-9].[shdb], z[0-9].[shdb]", assembly)
+    max_instr = re.findall(
+        r"max\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert (len(compare) > 1 and len(select) == len(compare)) or len(max_instr) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_min(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = T.min(A[v_i], B[v_i])
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    compare = re.findall(
+        r"cmgt\tp[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+    select = re.findall("sel\tz[0-9].[shdb], p[0-9], z[0-9].[shdb], z[0-9].[shdb]", assembly)
+    min_instr = re.findall(
+        r"min\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert (len(compare) > 1 and len(select) == len(compare)) or len(min_instr) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_div(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = tvm.tirx.div(A[v_i], B[v_i])
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"div\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) >= 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_mod(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = T.floormod(A[v_i], B[v_i])
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"mls\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 0
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_eq(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), "bool")
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] == B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"cm(p)?eq\tp[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    # The number of SVE compares depends on the LLVM cost model: LLVM <= 17
+    # interleaves the scalable loop by two, while LLVM 20 emits a fixed-width
+    # NEON main loop with a single predicated SVE epilogue.
+    assert len(matches) > 0
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_neq(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), "bool")
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] != B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"cm(p)?(gt|ne)\tp[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    # The number of SVE compares depends on the LLVM cost model: LLVM <= 17
+    # interleaves the scalable loop by two, while LLVM 20 emits a fixed-width
+    # NEON main loop with a single predicated SVE epilogue.
+    assert len(matches) > 0
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_or(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] | B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"orr\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_and(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] & B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"and\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_not(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = ~A[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"eor\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.xfail(
+    reason="Awaiting llvm support for gathered loads",
+    strict=True,
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_memcpy(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), "int32")
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(B[v_i], A[B[v_i]])
+                    T.writes(C[v_i])
+                    C[v_i] = A[B[v_i]]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+
+    assert len(loads) > 0
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 13,
+    reason="Function attribute vscale_range() is not supported in earlier versions of LLVM",
+)
+@pytest.mark.parametrize(
+    "mattr,expect_attr",
+    [
+        ("+neon", False),
+        ("+sve", True),
+        # Since LLVM 19, SVE/SVE2 are optional extensions of Armv9.0-A, so
+        # "+v9a" no longer implies "+sve" and no vscale_range() is added:
+        # https://releases.llvm.org/19.1.0/docs/ReleaseNotes.html#changes-to-the-aarch64-backend
+        ("+v9a", llvm_version_major() < 19),
+        ("+sme", True),
+    ],
+)
+def test_vscale_range_function_attribute(mattr, expect_attr):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": [mattr]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,))
+            C = T.match_buffer(var_C, (m,))
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] + T.float32(1)
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    # Check if the vscale_range() attribute exists
+    ll = f.inspect_source("ll")
+    attr = re.findall(r".*vscale_range\(\d+,\d+\)*.", ll)
+
+    if expect_attr:
+        assert len(attr) > 0, "Function attribute vscale_range() was not found in generated LLVM IR"
+    else:
+        assert len(attr) == 0, (
+            "Unexpected function attribute vscale_range() was found in generated LLVM IR"
+        )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_arm.py

@@ -0,0 +1,137 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import re
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_popcount():
+    target = {
+        "kind": "llvm",
+        "mtriple": "armv7l-none-linux-gnueabihf",
+        "mcpu": "cortex-a53",
+        "mattr": ["+neon"],
+    }
+
+    def check_correct_assembly(type, elements, counts):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((elements,), type), B: T.Buffer((elements,), type)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.vectorized(elements):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(elements, i)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = T.popcount(A[v_i])
+
+        f = tvm.tirx.build(Module, target=target)
+        # Verify we see the correct number of vpaddl and vcnt instructions in the assembly
+        assembly = f.inspect_source("asm")
+        matches = re.findall("vpaddl", assembly)
+        assert len(matches) == counts
+        matches = re.findall("vcnt", assembly)
+        assert len(matches) == 1
+
+    check_correct_assembly("uint16", 8, 1)
+    check_correct_assembly("uint16", 4, 1)
+    check_correct_assembly("uint32", 4, 2)
+    check_correct_assembly("uint32", 2, 2)
+    check_correct_assembly("uint64", 2, 3)
+
+
+def test_vmlal_s16():
+    target = {
+        "kind": "llvm",
+        "mtriple": "armv7l-none-linux-gnueabihf",
+        "mcpu": "cortex-a53",
+        "mattr": ["+neon"],
+    }
+
+    def check_correct_assembly(N):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(var_A: T.handle, var_B: T.handle, C: T.Buffer((N,), "int32")):
+                T.func_attr({"tirx.noalias": True})
+                K = T.int32()
+                A = T.match_buffer(var_A, (K, N), "int8")
+                B = T.match_buffer(var_B, (K, N), "int8")
+                for n in T.vectorized(N):
+                    for rv in range(K):
+                        with T.sblock("C"):
+                            v_n, v_rv = T.axis.remap("SR", [n, rv])
+                            T.reads(A[v_rv, v_n], B[v_rv, v_n])
+                            T.writes(C[v_n])
+                            with T.init():
+                                C[v_n] = 0
+                            C[v_n] = C[v_n] + T.Cast("int32", A[v_rv, v_n]) * T.Cast(
+                                "int32", B[v_rv, v_n]
+                            )
+
+        f = tvm.tirx.build(Module, target=target)
+
+        # Verify we see the correct number of vmlal.s16 instructions
+        assembly = f.inspect_source("asm")
+        matches = re.findall("vmlal.s16", assembly)
+        assert len(matches) == N // 4
+
+    check_correct_assembly(8)
+    check_correct_assembly(16)
+    check_correct_assembly(32)
+    check_correct_assembly(64)
+
+    def check_broadcast_correct_assembly(N):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(var_A: T.handle, var_B: T.handle, C: T.Buffer((N,), "int32")):
+                T.func_attr({"tirx.noalias": True})
+                K = T.int32()
+                A = T.match_buffer(var_A, (K, N), "int8")
+                B = T.match_buffer(var_B, (K,), "int8")
+                for n in T.vectorized(N):
+                    for rv in range(K):
+                        with T.sblock("C"):
+                            v_n, v_rv = T.axis.remap("SR", [n, rv])
+                            T.reads(A[v_rv, v_n], B[v_rv])
+                            T.writes(C[v_n])
+                            with T.init():
+                                C[v_n] = 0
+                            C[v_n] = C[v_n] + T.Cast("int32", A[v_rv, v_n]) * T.Cast(
+                                "int32", B[v_rv]
+                            )
+
+        f = tvm.tirx.build(Module, target=target)
+
+        # Verify we see the correct number of vmlal.s16 instructions
+        assembly = f.inspect_source("asm")
+        matches = re.findall("vmlal.s16", assembly)
+        assert len(matches) == N // 4
+
+    check_broadcast_correct_assembly(8)
+    check_broadcast_correct_assembly(16)
+    check_broadcast_correct_assembly(32)
+    check_broadcast_correct_assembly(64)
+
+
+if __name__ == "__main__":
+    test_popcount()
+    test_vmlal_s16()

tests/python/codegen/test_target_codegen_blob.py

@@ -0,0 +1,108 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F821
+
+import ctypes
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.support import cc, popen_pool, tar, utils
+
+
+@pytest.mark.gpu
+def test_cuda_multi_lib():
+    pytest.importorskip("cloudpickle")
+
+    # test combining two system lib together
+    # each contains a fatbin component in cuda
+    for device in ["llvm", "cuda"]:
+        if not tvm.testing.device_enabled(device):
+            print(f"skip because {device} is not enabled...")
+            return
+
+    @tvm.script.ir_module
+    class ModA:
+        I.module_attrs({"system_lib_prefix": "modA_"})
+
+        @T.prim_func(s_tir=True)
+        def my_inplace_update(x: T.Buffer((12), "float32")) -> None:
+            T.func_attr({"global_symbol": "modA_my_inplace_update"})
+            for bx in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for tx in T.thread_binding(T.int64(12), thread="threadIdx.x"):
+                    x[tx] = x[tx] + 1
+
+    @tvm.script.ir_module
+    class ModB:
+        I.module_attrs({"system_lib_prefix": "modB_"})
+
+        @T.prim_func(s_tir=True)
+        def my_inplace_update(x: T.Buffer((12), "float32")) -> None:
+            T.func_attr({"global_symbol": "modB_my_inplace_update"})
+            for bx in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for tx in T.thread_binding(T.int64(12), thread="threadIdx.x"):
+                    x[tx] = x[tx] + 2
+
+    temp = utils.tempdir()
+    target = tvm.target.Target("cuda", host="llvm")
+    libA = tvm.compile(ModA, target=target)
+    libB = tvm.compile(ModB, target=target)
+
+    pathA = temp.relpath("libA.tar")
+    pathB = temp.relpath("libB.tar")
+    pathAll = temp.relpath("libAll.a")
+
+    path_dso = temp.relpath("mylib.so")
+    libA.export_library(pathA, fcompile=tar.tar)
+    libB.export_library(pathB, fcompile=tar.tar)
+    cc.create_staticlib(pathAll, [pathA, pathB])
+    # package two static libs together
+    cc.create_shared(path_dso, ["-Wl,--whole-archive", pathAll, "-Wl,--no-whole-archive"])
+
+    def popen_check():
+        def run_and_check():
+            # Load dll, will trigger system library registration
+            ctypes.CDLL(path_dso)
+            # Load the system wide library
+            dev = tvm.cuda()
+            a_np = np.random.uniform(size=12).astype("float32")
+            a_nd = tvm.runtime.tensor(a_np, dev)
+            b_nd = tvm.runtime.tensor(a_np, dev)
+            syslibA = tvm.runtime.system_lib("modA_")
+            syslibB = tvm.runtime.system_lib("modB_")
+            # reload same lib twice
+            syslibA = tvm.runtime.system_lib("modA_")
+            syslibA["my_inplace_update"](a_nd)
+            syslibB["my_inplace_update"](b_nd)
+            np.testing.assert_equal(a_nd.numpy(), a_np + 1)
+            np.testing.assert_equal(b_nd.numpy(), a_np + 2)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    # system lib should be loaded in different process
+    worker = popen_pool.PopenWorker()
+    worker.send(popen_check)
+    worker.recv()
+
+
+if __name__ == "__main__":
+    test_synthetic()
+    test_cuda_multilib()

tests/python/codegen/test_target_codegen_bool.py

@@ -0,0 +1,109 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""codegen related to bool types"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("target", ["llvm", "cuda", "rocm", "vulkan", "metal", "opencl"])
+def test_cmp_load_store(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @I.ir_module(s_tir=True)
+    class GPUModule:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((32,), "float32"),
+            B: T.Buffer((32,), "float32"),
+            D: T.Buffer((32,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            C = T.sblock_alloc_buffer((32,), "bool")
+            for i0_0 in T.thread_binding(8, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(4, thread="blockIdx.x"):
+                    with T.sblock("C"):
+                        v_i0 = T.axis.spatial(32, i0_0 * 4 + i0_1)
+                        T.reads(B[v_i0], A[v_i0])
+                        T.writes(C[v_i0])
+                        C[v_i0] = B[v_i0] < A[v_i0]
+            for i0_0 in T.thread_binding(8, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(4, thread="blockIdx.x"):
+                    with T.sblock("D"):
+                        v_i0 = T.axis.spatial(32, i0_0 * 4 + i0_1)
+                        T.reads(C[v_i0], A[v_i0])
+                        T.writes(D[v_i0])
+                        D[v_i0] = T.Cast("float32", C[v_i0] and T.float32(1.0) < A[v_i0])
+
+    @I.ir_module(s_tir=True)
+    class CPUModule:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((32,), "float32"),
+            B: T.Buffer((32,), "float32"),
+            D: T.Buffer((32,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            C = T.sblock_alloc_buffer((32,), "bool")
+            for i0 in range(32):
+                with T.sblock("C"):
+                    v_i0 = T.axis.spatial(32, i0)
+                    T.reads(B[v_i0], A[v_i0])
+                    T.writes(C[v_i0])
+                    C[v_i0] = B[v_i0] < A[v_i0]
+            for i0 in range(32):
+                with T.sblock("D"):
+                    v_i0 = T.axis.spatial(32, i0)
+                    T.reads(C[v_i0], A[v_i0])
+                    T.writes(D[v_i0])
+                    D[v_i0] = T.Cast("float32", C[v_i0] and T.float32(1.0) < A[v_i0])
+
+    arr_size = 32
+    is_gpu = tvm.target.Target(target).kind.name != "llvm"
+    mod = GPUModule if is_gpu else CPUModule
+
+    f = tvm.compile(mod, target=target)
+
+    a_np = np.random.uniform(size=arr_size).astype("float32")
+    b_np = np.random.uniform(size=arr_size).astype("float32")
+
+    def run_and_check():
+        dev = tvm.device(target)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        d = tvm.runtime.tensor(np.zeros(arr_size, dtype="float32"), dev)
+        f(a, b, d)
+        np.testing.assert_equal(
+            d.numpy(),
+            np.logical_and(a_np > b_np, a_np > 1).astype("float32"),
+        )
+
+    if is_gpu:
+        tvm.testing.run_with_gpu_lock(run_and_check)
+    else:
+        run_and_check()
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_c_host.py

@@ -0,0 +1,249 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.support import utils
+
+
+def test_add():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_fadd(
+            A: T.Buffer((1024,), "float32"),
+            B: T.Buffer((1024,), "float32"),
+            C: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("C"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0], B[v_i0])
+                    T.writes(C[v_i0])
+                    C[v_i0] = A[v_i0] + B[v_i0]
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        fadd = m["test_fadd"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+        c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        fadd(a, b, c)
+        tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
+
+    check_c()
+
+
+def test_reinterpret():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_reinterpret(
+            A: T.Buffer((1024,), "int32"),
+            B: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("B"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0])
+                    T.writes(B[v_i0])
+                    B[v_i0] = T.reinterpret("float32", A[v_i0] + 2)
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        fadd = m["test_reinterpret"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.randint(-(2**30), 2**30, size=n).astype("int32"), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        fadd(a, b)
+        tvm.testing.assert_allclose(b.numpy(), (2 + a.numpy()).view("float32"))
+
+    check_c()
+
+
+def test_ceil():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_ceil(
+            A: T.Buffer((1024,), "float32"),
+            B: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("B"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0])
+                    T.writes(B[v_i0])
+                    B[v_i0] = T.ceil(A[v_i0])
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        fceil = m["test_ceil"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        fceil(a, b)
+        tvm.testing.assert_allclose(b.numpy(), (np.ceil(a.numpy()).view("float32")))
+
+    check_c()
+
+
+def test_floor():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_floor(
+            A: T.Buffer((1024,), "float32"),
+            B: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("B"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0])
+                    T.writes(B[v_i0])
+                    B[v_i0] = T.floor(A[v_i0])
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        ffloor = m["test_floor"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        ffloor(a, b)
+        tvm.testing.assert_allclose(b.numpy(), (np.floor(a.numpy()).view("float32")))
+
+    check_c()
+
+
+def test_round():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_round(
+            A: T.Buffer((1024,), "float32"),
+            B: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("B"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0])
+                    T.writes(B[v_i0])
+                    B[v_i0] = T.round(A[v_i0])
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        fround = m["test_round"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        fround(a, b)
+        tvm.testing.assert_allclose(b.numpy(), (np.round(a.numpy()).view("float32")))
+
+    check_c()
+
+
+def test_subroutine_call():
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, dtype="float32")):
+            Module.subroutine(A.data)
+
+        @T.prim_func(private=True, s_tir=True)
+        def subroutine(A_data: T.handle("float32")):
+            A = T.decl_buffer(1, dtype="float32", data=A_data)
+            A[0] = 42.0
+
+    built = tvm.tirx.build(Module, target="c")
+
+    source = built.inspect_source()
+    assert source.count("__tvm_ffi_main(void*") == 2, (
+        "Expected two occurrences, for forward-declaration and definition"
+    )
+    assert source.count("subroutine(float*") == 2, (
+        "Expected two occurrences, for forward-declaration and definition"
+    )
+    assert source.count("subroutine(") == 3, (
+        "Expected three occurrences, for forward-declaration, definition, and call from main."
+    )
+
+
+def test_workspace_allocation_cast():
+    @I.ir_module
+    class Module:
+        @T.prim_func
+        def main(A: T.Buffer((256,), "float32")):
+            workspace = T.alloc_buffer((256,), "float32", scope="global")
+            for i in range(256):
+                workspace[i] = A[i]
+            for i in range(256):
+                A[i] = workspace[i]
+
+    built = tvm.tirx.build(Module, target="c")
+    assert "((float*)TVMBackendAllocWorkspace(" in built.inspect_source()
+
+    temp = utils.tempdir()
+    built.export_library(temp.relpath("workspace.so"))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_cross_llvm.py

@@ -0,0 +1,103 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401, F841
+"""Test cross compilation"""
+
+import os
+import struct
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import rpc
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.support import cc, utils
+from tvm.testing import env
+
+
+@I.ir_module(s_tir=True)
+class AddModule:
+    @T.prim_func(s_tir=True)
+    def main(
+        A: T.Buffer((1024,), "float32"),
+        B: T.Buffer((1024,), "float32"),
+        C: T.Buffer((1024,), "float32"),
+    ):
+        T.func_attr({"tirx.noalias": True})
+        for i0_0 in T.parallel(256):
+            for i0_1 in T.vectorized(4):
+                with T.sblock("C"):
+                    v_i0 = T.axis.spatial(1024, i0_0 * 4 + i0_1)
+                    T.reads(A[v_i0], B[v_i0])
+                    T.writes(C[v_i0])
+                    C[v_i0] = A[v_i0] + B[v_i0]
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_llvm_add_pipeline():
+    nn = 1024
+
+    def verify_elf(path, e_machine):
+        with open(path, "rb") as fi:
+            arr = fi.read(20)
+            assert struct.unpack("ccc", arr[1:4]) == (b"E", b"L", b"F")
+            endian = struct.unpack("b", arr[0x5:0x6])[0]
+            endian = "<" if endian == 1 else ">"
+            assert struct.unpack(endian + "h", arr[0x12:0x14])[0] == e_machine
+
+    def build_arm():
+        target = {"kind": "llvm", "mtriple": "armv7-none-linux-gnueabihf"}
+        if not tvm.runtime.enabled("llvm"):
+            print(f"Skip because {target} is not enabled..")
+            return
+        temp = utils.tempdir()
+        f = tvm.tirx.build(AddModule, target=target)
+        path = temp.relpath("myadd.o")
+        f.write_to_file(path)
+        verify_elf(path, 0x28)
+        asm_path = temp.relpath("myadd.asm")
+        f.write_to_file(asm_path)
+        # Do a RPC verification, launch kernel on Arm Board if available.
+        host = os.environ.get("TVM_RPC_ARM_HOST", None)
+        remote = None
+        if host:
+            port = int(os.environ["TVM_RPC_ARM_PORT"])
+            try:
+                remote = rpc.connect(host, port)
+            except RuntimeError as e:
+                pass
+
+        if remote:
+            remote.upload(path)
+            farm = remote.load_module("myadd.o")
+            dev = remote.cpu(0)
+            n = nn
+            a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+            b = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+            c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+            farm(a, b, c)
+            tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
+            print("Verification finish on remote..")
+
+    build_arm()
+
+
+if __name__ == "__main__":
+    test_llvm_add_pipeline()

tests/python/codegen/test_target_codegen_cuda_fastmath.py

@@ -0,0 +1,306 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import re
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+import tvm.tirx as tirx
+from tvm.ir.module import IRModule
+from tvm.runtime.executable import Executable
+from tvm.script import tirx as T
+from tvm.support.nvcc import have_fp16
+from tvm.testing import env
+
+VECTOR_N_INPUTS = 8
+
+
+def make_prim_func(
+    name: str,
+    dtype: str,
+    num_inputs: int,
+    op: Callable[[tirx.Expr, ...], tirx.Expr],
+) -> tirx.PrimFunc:
+    """Make a primitive function that applies the given operation to the input buffer."""
+    if num_inputs == 1:
+
+        @T.prim_func
+        def kernel(
+            A: T.Buffer((VECTOR_N_INPUTS,), dtype),
+            B: T.Buffer((VECTOR_N_INPUTS,), dtype),
+        ):
+            T.func_attr({"global_symbol": name + "_kernel", "tirx.noalias": True})
+            for i in T.thread_binding(VECTOR_N_INPUTS, thread="threadIdx.x"):
+                B[i] = op(A[i])
+
+        return kernel
+    elif num_inputs == 2:
+
+        @T.prim_func
+        def kernel(
+            A: T.Buffer((VECTOR_N_INPUTS,), dtype),
+            E: T.Buffer((VECTOR_N_INPUTS,), dtype),
+            B: T.Buffer((VECTOR_N_INPUTS,), dtype),
+        ):
+            T.func_attr({"global_symbol": name + "_kernel", "tirx.noalias": True})
+            for i in T.thread_binding(VECTOR_N_INPUTS, thread="threadIdx.x"):
+                B[i] = op(A[i], E[i])
+
+        return kernel
+    else:
+        raise ValueError(f"Unsupported number of inputs: {num_inputs}")
+
+
+@dataclass(frozen=True)
+class MathCase:
+    name: str
+    op: Callable[[tirx.Expr, ...], tirx.Expr]
+    num_inputs: int
+    default_intrinsic_f16: str
+    default_intrinsic_bf16: str
+    default_intrinsic_f32: str
+    default_intrinsic_f64: str
+    fast_math_intrinsic_f32: str
+    np_ref: object
+    rtol: float = 1e-5
+    atol: float = 1e-6
+
+
+MATH_CASES = [
+    MathCase(
+        "exp_case",
+        T.exp,
+        1,
+        "hexp",
+        "hexp",
+        "expf",
+        "exp",
+        "__expf",
+        lambda x: np.exp(x),
+    ),
+    MathCase(
+        "exp10_case",
+        T.exp10,
+        1,
+        "hexp10",
+        "hexp10",
+        "exp10f",
+        "exp10",
+        "__exp10f",
+        lambda x: np.power(10.0, x),
+    ),
+    MathCase(
+        "log_case",
+        T.log,
+        1,
+        "hlog",
+        "hlog",
+        "logf",
+        "log",
+        "__logf",
+        lambda x: np.log(x),
+    ),
+    MathCase(
+        "log2_case",
+        T.log2,
+        1,
+        "hlog2",
+        "hlog2",
+        "log2f",
+        "log2",
+        "__log2f",
+        lambda x: np.log2(x),
+    ),
+    MathCase(
+        "log10_case",
+        T.log10,
+        1,
+        "hlog10",
+        "hlog10",
+        "log10f",
+        "log10",
+        "__log10f",
+        lambda x: np.log10(x),
+    ),
+    MathCase(
+        "tan_case",
+        T.tan,
+        1,
+        "htan",
+        "htan",
+        "tanf",
+        "tan",
+        "tanf",
+        lambda x: np.tan(x),
+    ),
+    MathCase(
+        "cos_case",
+        T.cos,
+        1,
+        "hcos",
+        "hcos",
+        "cosf",
+        "cos",
+        "__cosf",
+        lambda x: np.cos(x),
+    ),
+    MathCase(
+        "sin_case",
+        T.sin,
+        1,
+        "hsin",
+        "hsin",
+        "sinf",
+        "sin",
+        "__sinf",
+        lambda x: np.sin(x),
+    ),
+    MathCase(
+        "tanh_case",
+        T.tanh,
+        1,
+        "htanh",
+        "htanh",
+        "tanhf",
+        "tanh",
+        "__tanhf",
+        lambda x: np.tanh(x),
+    ),
+    MathCase(
+        "pow_case",
+        T.pow,
+        2,
+        "hpow",
+        "hpow",
+        "powf",
+        "pow",
+        "__powf",
+        lambda x, y: np.power(x, y),
+    ),
+]
+
+
+def make_mod(
+    dtype: str, case: MathCase, enable_fast_math: bool
+) -> tuple[tvm.target.Target, tvm.IRModule]:
+    """Make a module for the given dtype and case."""
+    target = tvm.target.Target("cuda")
+    prim_func = make_prim_func(case.name, dtype, case.num_inputs, case.op)
+    return target, tvm.IRModule.from_expr(prim_func.with_attr("target", target))
+
+
+def expected_intrinsic(dtype: str, case: MathCase, enable_fast_math: bool) -> str:
+    """Get the expected intrinsic for the given dtype and case."""
+    if dtype == "float16":
+        return case.default_intrinsic_f16
+    elif dtype == "bfloat16":
+        return case.default_intrinsic_bf16
+    elif dtype == "float32":
+        return case.fast_math_intrinsic_f32 if enable_fast_math else case.default_intrinsic_f32
+    elif dtype == "float64":
+        return case.default_intrinsic_f64
+    else:
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+
+def check_lowered_ir(
+    dtype: str, case: MathCase, enable_fast_math: bool
+) -> tuple[tvm.target.Target, IRModule]:
+    """Check the lowered IR for the given dtype and case."""
+    target, mod = make_mod(dtype, case, enable_fast_math)
+    with tvm.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
+        lowered_mod = tvm.tirx.transform.LowerIntrin()(mod)
+    script = lowered_mod.script(show_meta=False)
+    expected = expected_intrinsic(dtype, case, enable_fast_math)
+    assert re.search(rf"""["']{re.escape(expected)}["']""", script)
+    return target, lowered_mod
+
+
+def check_cuda_source(
+    target: tvm.target.Target,
+    mod: IRModule,
+    dtype: str,
+    case: MathCase,
+    enable_fast_math: bool,
+) -> Executable:
+    """Check the CUDA source for the given dtype and case."""
+    with tvm.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
+        executable = tvm.compile(mod, target=target)
+    source = executable.mod.imports[0].inspect_source()
+    expected = expected_intrinsic(dtype, case, enable_fast_math)
+    assert re.search(rf"(?<!_)\b{re.escape(expected)}\s*\(", source)
+    return executable
+
+
+def make_numpy_inputs(dtype: str, case: MathCase):
+    """Make the numpy inputs for the given dtype and case."""
+    lhs = np.array([0.25, 0.5, 1.0, 2.0, 4.0, 9.0, 16.0, 10.0], dtype=dtype)
+    if case.num_inputs == 1:
+        return [lhs]
+    elif case.num_inputs == 2:
+        rhs = np.array([2.0, 3.0, 0.5, 1.5, 0.25, 0.5, 2.0, 1.0], dtype=dtype)
+        return [lhs, rhs]
+    else:
+        raise ValueError(f"Unsupported number of inputs: {case.num_inputs}")
+
+
+def check_runtime(dtype: str, case: MathCase, executable: Executable):
+    """Check the runtime for the given dtype and case."""
+    np_inputs = make_numpy_inputs(dtype, case)
+    expected = case.np_ref(*[arr.astype(dtype) for arr in np_inputs]).astype(dtype)
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        tvm_inputs = [tvm.runtime.tensor(arr, device=dev) for arr in np_inputs]
+        output = tvm.runtime.empty((VECTOR_N_INPUTS,), dtype, dev)
+
+        executable(*tvm_inputs, output)
+
+        actual = output.numpy()
+        np.testing.assert_allclose(actual, expected, rtol=case.rtol, atol=case.atol)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("enable_fast_math", [False, True], ids=["default", "fast_math"])
+def test_cuda_math_intrinsic_lowering_pass_context(enable_fast_math):
+    check_lowered_ir("float32", MATH_CASES[0], enable_fast_math)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize(
+    "dtype",
+    ["float16", "bfloat16", "float32", "float64"],
+)
+@pytest.mark.parametrize("case", MATH_CASES, ids=lambda case: f"{case.name}")
+@pytest.mark.parametrize("enable_fast_math", [False, True], ids=["default", "fast_math"])
+def test_cuda_math_intrinsic_lowering_source_and_runtime(dtype, case, enable_fast_math):
+    if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
+        pytest.skip("GPU does not support float16")
+    if dtype == "bfloat16" and case.name.startswith("pow_"):
+        pytest.skip("pow_argnames=case is only supported for float")
+
+    target, lowered_mod = check_lowered_ir(dtype, case, enable_fast_math)
+    executable = check_cuda_source(target, lowered_mod, dtype, case, enable_fast_math)
+    check_runtime(dtype, case, executable)

tests/python/codegen/test_target_codegen_cuda_fp4.py

@@ -0,0 +1,278 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from itertools import product
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+try:
+    from ml_dtypes import float4_e2m1fn
+
+    ML_DTYPES_AVAILABLE = True
+except ImportError:
+    ML_DTYPES_AVAILABLE = False
+
+
+@pytest.mark.parametrize("promoted_dtype", ["float32x2", "float16x2"])
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_e2m1_vector_conversions(promoted_dtype):
+    native_dtype = "float4_e2m1fnx2"
+    vector_length = 64
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((vector_length,), native_dtype),
+            B: T.Buffer((vector_length,), native_dtype),
+            C: T.Buffer((vector_length,), native_dtype),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(vector_length // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(vector_length, i_0 * 32 + i_1)
+                        T.reads(A[v_i], B[v_i])
+                        T.writes(C[v_i])
+                        C[v_i] = T.Cast(
+                            native_dtype,
+                            T.Cast(promoted_dtype, A[v_i]) + T.Cast(promoted_dtype, B[v_i]),
+                        )
+
+    target = "cuda"
+    fadd = tvm.compile(Module, target=target)
+
+    if "x" in native_dtype:
+        lanes = int(native_dtype.split("x")[-1])
+    else:
+        lanes = 1
+
+    if "x" in promoted_dtype:
+        promoted_base_dtype = promoted_dtype.split("x")[0]
+    else:
+        promoted_base_dtype = promoted_dtype
+
+    np_shape = (vector_length, lanes) if lanes > 1 else (vector_length,)
+
+    # Create test data - either using ml_dtypes if available, or using int8 with valid FP4 values
+    if ML_DTYPES_AVAILABLE:
+        a_np = np.random.uniform(low=0, high=5, size=np_shape).astype(float4_e2m1fn)
+        b_np = np.random.uniform(low=0, high=5, size=np_shape).astype(float4_e2m1fn)
+    else:
+        # float4_e2m1fn possible values: [0, 0.5, 1, 1.5, 2, 3, 4, 6]
+        # We will create int8 arrays with valid FP4 bit patterns
+        valid_fp4_values = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]  # 4-bit values
+        a_np = np.random.choice(valid_fp4_values, size=np_shape).astype(np.int8)
+        b_np = np.random.choice(valid_fp4_values, size=np_shape).astype(np.int8)
+
+    def run_and_check():
+        dev = tvm.device(target, 0)
+        a = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev)
+        a.copyfrom(a_np)
+        b = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev)
+        b.copyfrom(b_np)
+        c = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev)
+        fadd(a, b, c)
+        # For the comparison, we will convert result to the promoted dtype and compare
+        # Note: When ml_dtypes is not available, we skip the numpy-level computation comparison
+        # and just verify that the CUDA kernel compiles and executes without error
+        c_result = c.numpy().astype(promoted_base_dtype)
+        if ML_DTYPES_AVAILABLE:
+            # Full comparison when ml_dtypes is available
+            expected = (a_np + b_np).astype(promoted_base_dtype)
+            tvm.testing.assert_allclose(c_result, expected)
+        else:
+            # When ml_dtypes is not available, we just verify the comparison ran successfully
+            # by checking that we got a result with the expected shape and dtype
+            assert c_result.shape == np_shape
+            assert c_result.dtype == promoted_base_dtype
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def _shuffle_reinterpret_module(n, num_blocks, vector_length, num_elem_per_storage):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((n // num_elem_per_storage,), "uint32"),
+            B: T.Buffer((n,), "float16"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    for i_2 in T.vectorized(vector_length):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(
+                                n, i_0 * 32 * vector_length + i_1 * vector_length + i_2
+                            )
+                            T.reads(A[v_i])
+                            T.writes(B[v_i])
+                            B[v_i] = T.Shuffle(
+                                [
+                                    T.reinterpret(
+                                        "float4_e2m1fnx2",
+                                        T.bitwise_and(
+                                            T.shift_right(
+                                                A[v_i // num_elem_per_storage],
+                                                ((v_i % num_elem_per_storage) // 2 * 4 * 2).astype(
+                                                    "uint32"
+                                                ),
+                                            ),
+                                            T.uint32((1 << (4 * 2)) - 1),
+                                        ).astype("uint8"),
+                                    ).astype("float16x2")
+                                ],
+                                indices=[v_i % 2],
+                            )
+
+    return Module
+
+
+def _scalar_reinterpret_module(n, num_blocks, vector_length, num_elem_per_storage):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((n // num_elem_per_storage,), "uint32"),
+            B: T.Buffer((n,), "float16"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    for i_2 in T.vectorized(vector_length):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(
+                                n, i_0 * 32 * vector_length + i_1 * vector_length + i_2
+                            )
+                            T.reads(A[v_i])
+                            T.writes(B[v_i])
+                            B[v_i] = T.reinterpret(
+                                "float4_e2m1fn",
+                                T.bitwise_and(
+                                    T.shift_right(
+                                        A[v_i // num_elem_per_storage],
+                                        (v_i % num_elem_per_storage * 4).astype("uint32"),
+                                    ),
+                                    T.uint32((1 << 4) - 1),
+                                ).astype("uint8"),
+                            ).astype("float16")
+
+    return Module
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_e2m1_dequantize():
+    n = 128
+
+    dev = tvm.device("cuda", 0)
+    target = tvm.target.Target.from_device(dev)
+    num_elem_per_storage = 32 // 4
+
+    # We only test the whether the code can be compiled.
+    for func_type, vector_length in product(["shuffle", "scalar"], [1, 2, 4]):
+        if func_type == "shuffle" and vector_length == 1:
+            # Vectorize is necessary for shuffle.
+            continue
+
+        num_blocks = n // (32 * vector_length)
+
+        if func_type == "shuffle":
+            mod = _shuffle_reinterpret_module(n, num_blocks, vector_length, num_elem_per_storage)
+        else:
+            mod = _scalar_reinterpret_module(n, num_blocks, vector_length, num_elem_per_storage)
+
+        tvm.compile(mod, target=target)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_e2m1_scalar_buffer_offset():
+    """Regression test: float4_e2m1fn scalar buffer access uses correct byte offset.
+
+    In CUDA sizeof(__nv_fp4_e2m1) = 1 byte, but fp4 data packs 2 elements per
+    byte.  GetBufferRef must emit ``index / 2`` so that the element index is
+    converted to the correct byte offset.  Without the fix the index was used
+    as-is, producing addresses 2x too large — reading garbage from out-of-bounds
+    memory instead of the correct fp4 value.
+
+    We verify by writing known fp4 values, casting each element to float16 on
+    the GPU, and checking the results match the expected fp4->fp16 conversion.
+    """
+    n = 128
+
+    @T.prim_func(s_tir=True)
+    def func(A_raw: T.Buffer((n // 2,), "uint8"), B: T.Buffer((n,), "float16")):
+        T.func_attr({"tir.noalias": True})
+        A = T.decl_buffer((n,), "float4_e2m1fn", data=A_raw.data)
+        for i in range(n):
+            with T.sblock("B"):
+                vi = T.axis.spatial(n, i)
+                T.reads(A[vi])
+                T.writes(B[vi])
+                B[vi] = T.Cast("float16", A[vi])
+
+    sch = tvm.s_tir.Schedule(func)
+    block = sch.get_sblock("B")
+    loops = sch.get_loops(block)
+    bx, tx = sch.split(loops[0], factors=[None, 32])
+    sch.bind(bx, "blockIdx.x")
+    sch.bind(tx, "threadIdx.x")
+
+    target = "cuda"
+    fadd = tvm.compile(sch.mod, target=target)
+
+    # float4_e2m1fn: 4-bit values 0..15, two packed per byte.
+    # Encoding (sign | exp1 | man1 man0):
+    #   0→0.0  1→0.5  2→1.0  3→1.5  4→2.0  5→3.0  6→4.0  7→6.0
+    #   8→-0.0 9→-0.5 10→-1.0 … 15→-6.0
+    fp4_to_fp16 = np.array(
+        [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
+        dtype=np.float16,
+    )
+
+    # Pack DIFFERENT fp4 values in low/high nibbles so the test verifies
+    # both byte offset (/2) AND correct nibble extraction (% 2 shift).
+    fp4_elements = np.array([i % 16 for i in range(n)], dtype=np.uint8)
+    packed = np.zeros(n // 2, dtype=np.uint8)
+    for i in range(0, n, 2):
+        packed[i // 2] = fp4_elements[i] | (fp4_elements[i + 1] << 4)
+
+    expected = fp4_to_fp16[fp4_elements]
+
+    def run_and_check():
+        dev = tvm.device(target, 0)
+        a = tvm.runtime.empty(shape=(n // 2,), dtype="uint8", device=dev)
+        a.copyfrom(packed)
+        b = tvm.runtime.empty(shape=(n,), dtype="float16", device=dev)
+        fadd(a, b)
+        tvm.testing.assert_allclose(b.numpy(), expected)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_device.py

@@ -0,0 +1,119 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
+def test_large_uint_imm():
+    value = (1 << 63) + 123
+    value_const = tvm.tirx.const(value, "uint64")
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((12,), "uint64")):
+            T.func_attr({"tirx.noalias": True})
+            for i0_0 in T.thread_binding(6, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(2, thread="threadIdx.x"):
+                    with T.sblock("A"):
+                        v_i0 = T.axis.spatial(12, i0_0 * 2 + i0_1)
+                        T.reads()
+                        T.writes(A[v_i0])
+                        A[v_i0] = value_const + T.uint64(3)
+
+    def check_target(target):
+        target_kind = target["kind"] if isinstance(target, dict) else target
+        if not tvm.testing.device_enabled(target_kind):
+            return
+        f = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target_kind, 0)
+            a = tvm.runtime.empty((12,), dtype="uint64", device=dev)
+            f(a)
+            assert a.numpy()[0] == value + 3
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_target("cuda")
+    check_target({"kind": "vulkan", "from_device": 0})
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
+def test_add_pipeline():
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, B: T.Buffer((), "float32"), var_D: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            n = T.int32()
+            A = T.match_buffer(var_A, (n,))
+            D = T.match_buffer(var_D, (n,))
+            C = T.sblock_alloc_buffer((n,))
+            for i0_0 in T.thread_binding((n + 255) // 256, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(256, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i0 = T.axis.spatial(n, i0_0 * 256 + i0_1)
+                        T.where(i0_0 * 256 + i0_1 < n)
+                        T.reads(A[v_i0], B[()])
+                        T.writes(C[v_i0])
+                        C[v_i0] = A[v_i0] + B[()]
+            for i0_0 in T.thread_binding((n + 255) // 256, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(256, thread="threadIdx.x"):
+                    with T.sblock("D"):
+                        v_i0 = T.axis.spatial(n, i0_0 * 256 + i0_1)
+                        T.where(i0_0 * 256 + i0_1 < n)
+                        T.reads(C[v_i0])
+                        T.writes(D[v_i0])
+                        D[v_i0] = C[v_i0] + T.float32(1.0)
+
+    def check_target(device, host):
+        if not tvm.testing.device_enabled(device) or not tvm.testing.device_enabled(host):
+            return
+        target = tvm.target.Target(device, host)
+        mhost = tvm.tirx.build(Module, target=target)
+        f = mhost.main
+        n = 1027
+
+        def run_and_check():
+            dev = tvm.device(device, 0)
+            a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+            b = tvm.runtime.tensor(np.random.uniform(size=()).astype("float32"), dev)
+            d = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+            f(a, b, d)
+            tvm.testing.assert_allclose(d.numpy(), a.numpy() + b.numpy() + 1)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_target("cuda", host="llvm")
+    # check_target("nvptx", host="llvm")  # nvptx kernel entry-point lookup not wired here
+    check_target("vulkan", host="llvm")
+    check_target("rocm", host="llvm")
+
+
+if __name__ == "__main__":
+    test_large_uint_imm()
+    test_add_pipeline()

tests/python/codegen/test_target_codegen_extern.py

@@ -0,0 +1,109 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+@pytest.mark.gpu
+def test_add_pipeline():
+    """Test extern-style add pipeline with vectorized operations."""
+    nn = 64
+    max_threads = 4
+
+    # CPU version: serial loop with vectorized operations
+    @I.ir_module
+    class ModuleCPU:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((64,), "float32"), C: T.Buffer((64,), "float32")):
+            for i in T.serial((64 + 1) // 2):
+                C[T.Ramp(i * 2, 1, 2)] = A[T.Ramp(i * 2, 1, 2)] + T.Broadcast(T.float32(1), 2)
+
+    # GPU version: thread bindings with vectorized operations
+    @I.ir_module
+    class ModuleGPU:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((64,), "float32"), C: T.Buffer((64,), "float32")):
+            bx = T.launch_thread("blockIdx.x", (64 + 4 - 1) // 4)
+            tx = T.launch_thread("threadIdx.x", 4)
+            idx = bx * 4 + tx
+            if T.likely(idx < 64):
+                C[T.Ramp(idx * 2, 1, 2)] = A[T.Ramp(idx * 2, 1, 2)] + T.Broadcast(T.float32(1), 2)
+
+    def check_target(target):
+        if not tvm.testing.device_enabled(target):
+            return
+        mod = ModuleGPU if target in ["opencl", "cuda"] else ModuleCPU
+        # build and invoke the kernel.
+        f = tvm.compile(mod, target=target)
+        n = nn
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+            c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+            f(a, c)
+            tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)
+
+        if target == "llvm":
+            run_and_check()
+        else:
+            tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_target("llvm")
+    check_target("opencl")
+    check_target("cuda")
+
+
+def test_pack_buffer_simple():
+    """Test call_packed with buffer arguments."""
+    nn = 1024
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1024,), "float32"), C: T.Buffer((1024,), "float32")):
+            T.evaluate(T.call_packed("my_extern_array_func1", A, C))
+
+    @tvm.register_global_func
+    def my_extern_array_func1(aa, bb):
+        aa.copyto(bb)
+
+    def check_target(target):
+        if not tvm.testing.device_enabled(target):
+            return
+        # build and invoke the kernel.
+        f = tvm.compile(Module, target=target)
+        dev = tvm.cpu(0)
+        # launch the kernel.
+        n = nn
+        a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+        c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+
+        f(a, c)
+        tvm.testing.assert_allclose(c.numpy(), a.numpy())
+
+    check_target("llvm")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_gpu_common.py

@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from functools import partial
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+        pytest.param({"kind": "vulkan", "supports_int64": True}, marks=pytest.mark.gpu),
+        pytest.param("opencl", marks=pytest.mark.gpu),
+    ],
+)
+@pytest.mark.parametrize("dtype", ["int32", "uint32", "int64", "uint64"])
+def test_int_intrin(target, dtype):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    test_funcs = [
+        (T.clz, lambda x, dtype: int(dtype[-2:]) - (len(bin(x)) - 2)),
+    ]
+
+    for tvm_intrin, np_func in test_funcs:
+        n = 128
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(
+                A: T.Buffer((n,), dtype),
+                B: T.Buffer((n,), dtype),
+            ):
+                T.func_attr({"tirx.noalias": True})
+                for i0 in T.thread_binding(n, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i0 = T.axis.spatial(n, i0)
+                        T.reads(A[v_i0])
+                        T.writes(B[v_i0])
+                        B[v_i0] = tvm_intrin(A[v_i0])
+
+        f = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target["kind"] if isinstance(target, dict) else target)
+            a = tvm.runtime.tensor(np.random.randint(0, 100000, size=n).astype(dtype), dev)
+            b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev)
+            f(a, b)
+            ref = np.vectorize(partial(np_func, dtype=dtype))(a.numpy())
+            tvm.testing.assert_allclose(b.numpy(), ref)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_hexagon.py

@@ -0,0 +1,122 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import re
+
+import pytest
+
+import tvm
+import tvm.contrib.hexagon as hexagon
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.fixture(autouse=True)
+def register_linker():
+    original_linker = hexagon.hexagon_link()
+    # Register a phony linker, so that we can test codegen without a Hexagon toolchain.
+    hexagon.register_linker(lambda: "/bin/true")
+    yield None
+    # Restore registration.
+    hexagon.register_linker(original_linker)
+
+
+@pytest.mark.skipif(not env.has_hexagon(), reason="need hexagon")
+def test_basic():
+    target = tvm.target.Target("qcom/hexagon-v66")
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            C: T.Buffer((128,), "uint8"),
+            A: T.Buffer((128,), "uint8"),
+            A_1: T.Buffer((128,), "uint8"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i in range(128):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(128, i)
+                    T.reads(A[v_i], A_1[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] + A_1[v_i]
+
+    hexm = tvm.compile(Module, target=tvm.target.Target(target, target))
+    asm = hexm.inspect_source("s")
+    vadds = re.findall(r"v[0-9]+.b = vadd\(v[0-9]+.b,v[0-9]+.b\)", asm)
+    assert vadds  # Check that it's non-empty
+
+
+@pytest.mark.skipif(not env.has_hexagon(), reason="need hexagon")
+def test_llvm_target_features():
+    target = tvm.target.Target("qcom/hexagon-v66")
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def add_one(C: T.Buffer((128,), "int32"), A: T.Buffer((128,), "uint8")):
+            T.func_attr({"tirx.noalias": True})
+            for i in range(128):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(128, i)
+                    T.reads(A[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = T.Cast("int32", A[v_i]) + 1
+
+    m = tvm.compile(Module, target=tvm.target.Target(target, target))
+    llvm_ir = m.inspect_source("ll")
+    # Make sure we find +hvx-length128b in "attributes".
+    fs = re.findall(r"attributes.*\+hvx-length128b", llvm_ir)
+    assert fs  # Check that it's non-empty
+
+
+@pytest.mark.skipif(not env.has_hexagon(), reason="need hexagon")
+def test_llvm_options():
+    target = tvm.target.Target(
+        {
+            "kind": "hexagon",
+            "mtriple": "hexagon",
+            "mcpu": "hexagonv66",
+            "mattr": ["+hvxv66", "+hvx-length128b"],
+            "num-cores": 4,
+            "vtcm-capacity": 262144,
+            "llvm-options": ["-hexagon-noopt"],
+        }
+    )
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(compute: T.Buffer((10,), "int32")):
+            T.func_attr({"tirx.noalias": True})
+            for _ in range(10):
+                with T.sblock("compute"):
+                    v__ = T.axis.spatial(10, _)
+                    T.reads()
+                    T.writes(compute[v__])
+                    compute[v__] = 0
+
+    # Check that BuildHexagon hasn't crashed because of target attribute
+    # type mismatch.
+    tvm.compile(Module, target=tvm.target.Target(target, target))
+    assert re.search("-hexagon-noopt", str(target))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_llvm_vla.py

@@ -0,0 +1,212 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Codegen tests for VLA extensions
+"""
+
+import re
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.target.codegen import llvm_version_major
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_codegen_vscale(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    vscale = tvm.tirx.vscale()
+
+    @T.prim_func(s_tir=True)
+    def main(A: T.Buffer((5,), "int32")):
+        for i in range(5):
+            A[i] = 2 * vscale
+
+    with tvm.target.Target(target):
+        build_mod = tvm.tirx.build(main)
+
+    llvm = build_mod.inspect_source()
+    assert re.findall(r"llvm.vscale.i32", llvm), "No vscale in generated LLVM."
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_scalable_buffer_load_store(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def my_func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (128,), "float32")
+        B = T.match_buffer(b, (128,), "float32")
+        T.func_attr({"global_symbol": "my_module", "tirx.noalias": True})
+        B[T.ramp(0, 1, 4 * T.vscale())] = A[T.ramp(0, 1, 4 * T.vscale())]
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.build(my_func)
+
+    llvm = mod.inspect_source("ll")
+    assert re.findall(r"load <vscale x 4 x float>", llvm), "No scalable load in generated LLVM."
+    assert re.findall(r" store <vscale x 4 x float>", llvm), "No scalable store in generated LLVM."
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_scalable_broadcast(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def my_func(a: T.handle):
+        A = T.match_buffer(a, (128,), "float32")
+        T.func_attr({"global_symbol": "my_module", "tirx.noalias": True})
+        A[T.ramp(0, 1, 4 * T.vscale())] = T.broadcast(1, 4 * T.vscale())
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.build(my_func)
+
+    llvm = mod.inspect_source("ll")
+    # Older LLVM versions print the broadcast as a shufflevector of an insertelement,
+    # newer ones print it as a splat constant.
+    assert (
+        "shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float>" in llvm
+        or "store <vscale x 4 x float> splat (float 1.000000e+00)" in llvm
+    ), "No scalable broadcast in generated LLVM."
+    assert re.findall(r" store <vscale x 4 x float>", llvm), "No scalable store in generated LLVM."
+
+
+@pytest.mark.skip(
+    reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_get_active_lane_mask(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle):
+        A = T.match_buffer(a, (30,), "int1")
+        for i in range(T.ceildiv(30, T.vscale() * 4)):
+            A[i : i + T.vscale() * 4] = T.get_active_lane_mask("uint1xvscalex4", i, 30)
+
+    with tvm.target.Target(target):
+        out = tvm.tirx.build(before)
+
+    ll = out.inspect_source("ll")
+    assert "get.active.lane.mask" in ll
+
+
+@pytest.mark.skip(
+    reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_predicated_scalable_buffer(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in T.serial(T.ceildiv(16, 4 * T.vscale())):
+            for i_1 in T.vectorized(4 * T.vscale()):
+                if i_0 * 4 * T.vscale() + i_1 < 14:
+                    B[i_0 * 4 * T.vscale() + i_1] = A[i_0 * 4 * T.vscale() + i_1] + 1.0
+
+    with tvm.target.Target(target):
+        out = tvm.tirx.build(before)
+
+    ll = out.inspect_source("ll")
+    assert "get.active.lane.mask" in ll
+    assert "llvm.masked.load" in ll
+    assert "llvm.masked.store" in ll
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_metal.py

@@ -0,0 +1,353 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_metal_inf_nan():
+    target = "metal"
+
+    def check_inf_nan(n, value, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(
+                A: T.Buffer((1,), dtype),
+                C: T.Buffer((1,), dtype),
+            ):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads()
+                        T.writes(C[v_i])
+                        C[v_i] = T.Cast(dtype, value)
+
+        fun = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_inf_nan(1, -float("inf"), "float32")
+    check_inf_nan(1, -float("inf"), "float16")
+    check_inf_nan(1, float("inf"), "float32")
+    check_inf_nan(1, float("inf"), "float16")
+    check_inf_nan(1, float("nan"), "float32")
+    check_inf_nan(1, float("nan"), "float16")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_unaligned_vectorize():
+    @tvm.script.ir_module
+    class IRModule:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((2, 3), "float32"), B: T.Buffer((6,), "float32")):
+            T.func_attr({"global_symbol": "main"})
+            for i0_1 in T.thread_binding(3, thread="threadIdx.x"):
+                for i0_0 in T.vectorized(2):
+                    with T.sblock("block"):
+                        vi0 = T.axis.spatial(6, i0_0 * 3 + i0_1)
+                        B[vi0] = A[vi0 // 3, vi0 % 3]
+
+    target = "metal"
+    a = (np.arange(6).reshape(2, 3)).astype("float32")
+    f = tvm.compile(IRModule, target=target)
+
+    def run_and_check():
+        dev = tvm.metal()
+        a_nd = tvm.runtime.tensor(a, dev)
+        b_nd = tvm.runtime.empty((6,), "float32", dev)
+        f(a_nd, b_nd)
+        tvm.testing.assert_allclose(b_nd.numpy(), a.reshape(6), atol=1e-5, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_metal_erf():
+    target = "metal"
+
+    def check_erf(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(
+                A: T.Buffer((1,), dtype),
+                C: T.Buffer((1,), dtype),
+            ):
+                T.func_attr({"tirx.noalias": True})
+                for i0 in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i0 = T.axis.spatial(1, i0)
+                        T.reads(A[v_i0])
+                        T.writes(C[v_i0])
+                        C[v_i0] = T.erf(A[v_i0])
+
+        fun = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_erf(1, "float32")
+    check_erf(1, "float16")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_ramp():
+    target = "metal"
+
+    @tvm.script.ir_module
+    class IRModule:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1, 2), "int32")):
+            T.func_attr({"global_symbol": "main"})
+            for i in T.thread_binding(1, thread="threadIdx.x"):
+                with T.sblock("block"):
+                    tx = T.axis.spatial(1, i)
+                    r = T.ramp(tx, 3, 2)
+                    A[0, T.ramp(0, 1, 2)] = r
+
+    f = tvm.compile(IRModule, target=target)
+
+    def run_and_check():
+        dev = tvm.metal()
+        a_nd = tvm.runtime.empty((1, 2), "int32", dev)
+        f(a_nd)
+        assert tuple(a_nd.numpy()[0, :]) == (0, 3)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_select_vectorize():
+    @tvm.script.ir_module
+    class IRModule:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((6), "float32"), B: T.Buffer((6,), "float32")):
+            T.func_attr({"global_symbol": "main"})
+            for i0_1 in T.thread_binding(3, thread="threadIdx.x"):
+                for i0_0 in T.vectorized(2):
+                    with T.sblock("block"):
+                        vi0 = T.axis.spatial(6, i0_0 * 3 + i0_1)
+                        B[vi0] = T.Select((vi0 % 2) == 0, A[vi0], T.float32(0))
+
+    target = "metal"
+    a = np.arange(6).astype("float32")
+    f = tvm.compile(IRModule, target=target)
+    a.reshape(3, 2)[:, 1] = 0
+
+    def run_and_check():
+        dev = tvm.metal()
+        a_nd = tvm.runtime.tensor(a, dev)
+        b_nd = tvm.runtime.empty((6,), "float32", dev)
+        f(a_nd, b_nd)
+        tvm.testing.assert_allclose(b_nd.numpy(), a, atol=1e-5, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_vectorized_uint8():
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((16), "uint8"), B: T.Buffer((16), "float32")):
+        for i in T.thread_binding(4, thread="threadIdx.x"):
+            for j in T.vectorized(4):
+                with T.sblock("block"):
+                    vi = T.axis.spatial(16, i * 4 + j)
+                    B[vi] = T.Cast("float32", A[vi])
+
+    a = np.arange(16).astype("uint8")
+    f = tvm.compile(func, target="metal")
+
+    def run_and_check():
+        dev = tvm.metal()
+        a_nd = tvm.runtime.tensor(a, dev)
+        b_nd = tvm.runtime.empty((16,), "float32", dev)
+        f(a_nd, b_nd)
+        tvm.testing.assert_allclose(b_nd.numpy(), a.astype("float32"), atol=1e-5, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_func_with_trailing_pod_params():
+    from tvm.support import xcode  # pylint: disable=import-outside-toplevel
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((16), "float32"), B: T.Buffer((16), "float32"), x: T.float32):
+        for i in T.thread_binding(16, thread="threadIdx.x"):
+            with T.sblock("block"):
+                vi = T.axis.spatial(16, i)
+                B[vi] = A[vi] + x
+
+    @tvm.register_global_func("tvm_callback_metal_compile")
+    def compile_metal(src, target):
+        return xcode.compile_metal(src)
+
+    mod = tvm.IRModule({"main": func})
+
+    f = tvm.tirx.build(mod, target="metal")
+    src: str = f.imports[0].inspect_source()
+    occurrences = src.count("struct func_kernel_args_t")
+    assert occurrences == 1, occurrences
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_metal_compile_callback_source_passthrough():
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    seen = {}
+
+    def inspect_callback(src, target):
+        # Pure inspection callback: capture the source, return it untouched and
+        # declare it is still textual MSL so it is compiled at load time.
+        seen["src"] = src
+        return (src, "metal")
+
+    tvm.register_global_func("tvm_callback_metal_compile", inspect_callback, override=True)
+    try:
+        f = tvm.compile(Module, target="metal")
+        dev = tvm.metal()
+        a = np.random.rand(n).astype("float32")
+        a_nd = tvm.runtime.tensor(a, dev)
+        b_nd = tvm.runtime.empty((n,), "float32", dev)
+        f(a_nd, b_nd)
+        dev.sync()
+    finally:
+        tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile")
+
+    assert "src" in seen and len(seen["src"]) > 0
+    tvm.testing.assert_allclose(b_nd.numpy(), a + 1.0, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_metal_compile_callback_mixed_formats_rejected():
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((n,), "float32"),
+            B: T.Buffer((n,), "float32"),
+            C: T.Buffer((n,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            # Two independent thread-bound regions -> two device kernels, so the
+            # compile callback is invoked twice within one module.
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+            for j_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for j_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_j = T.axis.spatial(n, j_0 * 32 + j_1)
+                        T.reads(A[v_j])
+                        T.writes(C[v_j])
+                        C[v_j] = A[v_j] + 2.0
+
+    calls = {"n": 0}
+
+    def mixed_callback(src, target):
+        calls["n"] += 1
+        if calls["n"] == 1:
+            # Treated as a compiled metallib payload.
+            return src
+        # Second kernel declares textual MSL, contradicting the metallib above.
+        return (src, "metal")
+
+    tvm.register_global_func("tvm_callback_metal_compile", mixed_callback, override=True)
+    try:
+        with pytest.raises(Exception, match="inconsistent formats"):
+            tvm.compile(Module, target="metal")
+    finally:
+        tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_export_load_with_fallback(monkeypatch, tmp_path):
+    """Force the codegen wrapper into the fallback branch, then export."""
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+    host_lib = tvm.compile(Module, target="metal")
+    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
+
+    lib_path = str(tmp_path / "lib.so")
+    host_lib.export_library(lib_path)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_opencl.py

@@ -0,0 +1,334 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501
+import re
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+target = "opencl"
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_opencl_ternary_expression():
+    def check_if_then_else(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads(A[0])
+                        T.writes(C[v_i])
+                        C[v_i] = T.max(
+                            T.Cast(dtype, 2),
+                            T.if_then_else(
+                                0 < T.Cast("int32", A[0]),
+                                T.Cast(dtype, 1),
+                                T.Cast(dtype, 3),
+                            ),
+                        )
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    def check_select(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads(A[0])
+                        T.writes(C[v_i])
+                        C[v_i] = T.max(
+                            T.Cast(dtype, 2),
+                            T.Select(
+                                0 < T.Cast("int32", A[0]),
+                                T.Cast(dtype, 1),
+                                T.Cast(dtype, 3),
+                            ),
+                        )
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_if_then_else(1, "int8")
+    check_if_then_else(1, "uint8")
+    check_if_then_else(1, "int16")
+    check_if_then_else(1, "uint16")
+    check_select(1, "int8")
+    check_select(1, "uint8")
+    check_select(1, "int16")
+    check_select(1, "uint16")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_opencl_inf_nan():
+    def check_inf_nan(n, value, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads()
+                        T.writes(C[v_i])
+                        C[v_i] = T.Cast(dtype, value)
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_inf_nan(1, -float("inf"), "float32")
+    check_inf_nan(1, -float("inf"), "float64")
+    check_inf_nan(1, float("inf"), "float32")
+    check_inf_nan(1, float("inf"), "float64")
+    check_inf_nan(1, float("nan"), "float32")
+    check_inf_nan(1, float("nan"), "float64")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_opencl_max():
+    def check_max(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads(A[0])
+                        T.writes(C[v_i])
+                        C[v_i] = T.max(A[0] + T.Cast(dtype, 1), T.Cast(dtype, 0))
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_max(1, "int8")
+    check_max(1, "uint8")
+    check_max(1, "int16")
+    check_max(1, "uint16")
+    check_max(1, "float32")
+    check_max(1, "float64")
+
+
+def test_opencl_erf():
+    def check_erf(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i0 in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i0 = T.axis.spatial(1, i0)
+                        T.reads(A[v_i0])
+                        T.writes(C[v_i0])
+                        C[v_i0] = T.erf(A[v_i0])
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        source_str = fun.imports[0].inspect_source()
+        matches = re.findall("erf", source_str)
+        error_matches = re.findall("erff", source_str)
+        assert len(matches) == 1 and len(error_matches) == 0
+
+    check_erf(1, "float32")
+    check_erf(1, "float64")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_opencl_type_casting():
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(C: T.Buffer((32,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(8, thread="threadIdx.x"):
+                for i_1 in T.vectorized(4):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(32, i_0 * 4 + i_1)
+                        T.reads()
+                        T.writes(C[v_i])
+                        C[v_i] = T.Select(
+                            v_i // 4 == 3 and v_i % 3 == 1, T.float32(1.0), T.float32(0.0)
+                        )
+
+    def check_type_casting(n, dtype):
+        fun = tvm.tirx.build(Module, target=target)
+        assembly = fun.imports[0].inspect_source()
+        lcond = "convert_int4(((convert_uint4(((uint4)(((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3)))))"
+        rcond = "(convert_uint4(((((int4)(((convert_int(get_local_id(0))))+(1*0), ((convert_int(get_local_id(0))))+(1*1), ((convert_int(get_local_id(0))))+(1*2), ((convert_int(get_local_id(0))))+(1*3))) % ((int4)(3, 3, 3, 3))) == ((int4)(1, 1, 1, 1))))))))"
+        pattern_cond = f"({lcond} && {rcond})"
+        assert assembly.count(pattern_cond) != 0
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_type_casting(32, "float32")
+    # fp16 is not yet supported in ci
+    # check_type_casting(dev, 16, "float16")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("opencl", marks=pytest.mark.gpu),
+        pytest.param({"kind": "opencl", "device": "adreno"}, marks=pytest.mark.gpu),
+    ],
+)
+def test_opencl_ceil_log2(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    def _check(target, n, dtype):
+        target_obj = tvm.target.Target(target)
+        is_adreno = "adreno" in target_obj.attrs.get("device", "")
+        inter_dtype = "float32" if is_adreno else "float64"
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(C: T.Buffer((n,), "int32")):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(n, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(n, i)
+                        T.reads()
+                        T.writes(C[v_i])
+                        C[v_i] = T.Cast("int32", T.ceil(T.log2(T.Cast(inter_dtype, v_i))))
+
+        fun = tvm.tirx.build(Module, target=target)
+        assembly = fun.imports[0].inspect_source()
+        if is_adreno:
+            pattern = "convert_float"
+        else:
+            pattern = "convert_double"
+        assert assembly.count(pattern) != 0
+
+    _check(target, 32, "float32")
+
+
+def _get_maximum_kernel_args(source):
+    def get_kernel_args(source):
+        import re
+
+        p = re.tirx.build(r"__kernel void .+\((.*)\)")
+        args = p.findall(source)
+        return args
+
+    args = get_kernel_args(source)
+    max_args = len(args[0].split(","))
+    for arg_line in args:
+        max_args = max(max_args, len(arg_line.split(",")))
+    return max_args
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_export_load_with_fallback(monkeypatch, tmp_path):
+    """Force the codegen wrapper into the fallback branch, then export+load+run."""
+    import numpy as np
+
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+    host_lib = tvm.compile(Module, target=target)
+    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
+
+    lib_path = str(tmp_path / "lib.so")
+    host_lib.export_library(lib_path)
+    reloaded = tvm.runtime.load_module(lib_path)
+
+    a_np = np.random.uniform(size=(n,)).astype("float32")
+    b_np = np.zeros((n,), dtype="float32")
+
+    def run_and_check():
+        dev = tvm.device(target, 0)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        reloaded["main"](a, b)
+        np.testing.assert_allclose(b.numpy(), a_np + 1.0, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_riscv.py

@@ -0,0 +1,199 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501, F841
+
+import re
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.target.codegen import target_has_features
+from tvm.testing import env
+
+
+@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14")
+@pytest.mark.parametrize(
+    "target",
+    [
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv32-linux-gnu",
+            "mcpu": "generic-rv32",
+            "mattr": ["+i", "+m"],
+        },
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv32-linux-gnu",
+            "mcpu": "generic-rv32",
+            "mattr": ["+i", "+m", "+v"],
+        },
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m"],
+        },
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_rvv(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    def check_rvv_presence(N, extent):
+        @T.prim_func(s_tir=True)
+        def load_vec(A: T.Buffer((N,), "int8")):
+            for j in T.vectorized(0, extent):
+                A[j] = 1
+
+        f = tvm.tirx.build(load_vec, target)
+        # Check RVV `vsetvli` prensence
+        assembly = f.inspect_source("asm")
+        if target_has_features("v"):
+            assert "vsetvli" in assembly
+        else:
+            assert "vsetvli" not in assembly
+
+    with tvm.target.Target(target):
+        check_rvv_presence(16, 32)
+
+
+@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14")
+@pytest.mark.parametrize(
+    "target",
+    [
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv32-linux-gnu",
+            "mcpu": "generic-rv32",
+            "mattr": ["+i", "+m", "+v"],
+        },
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_rvv_vscale_llvm_dbginfo(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def rvv_with_vscale(A_handle: T.handle, B_handle: T.handle, C_handle: T.handle):
+        A = T.match_buffer(A_handle, (8,), dtype="float32", align=4, offset_factor=1)
+        B = T.match_buffer(B_handle, (4, 8), dtype="float32", align=4, offset_factor=1, strides=[8, 1])
+        C = T.match_buffer(C_handle, (4,), dtype="float32", align=4, offset_factor=1)
+        with T.sblock("root"):
+            T.reads(A[0:8], B[0:4, 0:8])
+            zero = T.call_llvm_intrin("float32xvscalex2", "llvm.riscv.vfmv.v.f", T.Broadcast(T.float32(0.0), T.vscale() * 2), C[0], T.uint64(1))
+            vec_A = T.call_llvm_intrin("float32xvscalex4", "llvm.riscv.vle", T.Broadcast(T.float32(0.0), T.vscale() * 4), T.tvm_access_ptr(T.type_annotation("float32"), A.data, 0, 8, 1), T.int64(8))
+            vec_B = T.call_llvm_intrin("float32xvscalex4", "llvm.riscv.vle", T.Broadcast(T.float32(0.0), T.vscale() * 4), T.tvm_access_ptr(T.type_annotation("float32"), B.data, 0 * 8, 8, 1), T.int64(8))
+            prod = T.call_llvm_intrin("float32xvscalex4", "llvm.riscv.vfmul", T.Broadcast(T.float32(0.0), T.vscale() * 4), vec_A, vec_B, T.uint64(7), T.uint64(8))
+            redsum = T.call_llvm_intrin("float32xvscalex2", "llvm.riscv.vfredusum", T.Broadcast(T.float32(0.0), T.vscale() * 2), prod, zero, T.uint64(7), T.uint64(8))
+    # fmt: on
+
+    # tvm.error.InternalError: Can't fetch the lanes of a scalable vector at a compile time.
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(rvv_with_vscale, target)
+
+
+@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14")
+def test_rvv_fixed_width_vectorized_loop_uses_scalable_chunks():
+    @T.prim_func(s_tir=True)
+    def fixed16_negative(
+        A: T.Buffer((14, 23, 67, 99), "float32"),
+        B: T.Buffer((14, 23, 67, 99), "float32"),
+    ):
+        for n, c, h, wo in T.grid(14, 23, 67, 7):
+            for wi in T.vectorized(0, 16):
+                if wo * 16 + wi < 99:
+                    B[n, c, h, wo * 16 + wi] = T.float32(0) - A[n, c, h, wo * 16 + wi]
+
+    @T.prim_func(s_tir=True)
+    def fixed16_negative_int64(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
+        for wi in T.vectorized(T.int64(0), T.int64(16)):
+            B[wi] = T.float32(0) - A[wi]
+
+    target = tvm.target.Target(
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        }
+    )
+
+    def check_codegen(func):
+        with target:
+            f = tvm.tirx.build(func, target)
+
+        assembly = f.inspect_source("asm")
+        assert "vle32.v" in assembly
+        assert "vse32.v" in assembly
+        assert not re.search(r"\bflw\b", assembly)
+        assert not re.search(r"\bfsub\.s\b", assembly)
+        assert not re.search(r"\bfsw\b", assembly)
+
+    check_codegen(fixed16_negative)
+    check_codegen(fixed16_negative_int64)
+
+
+@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14")
+def test_rvv_scalable_ramp_expression():
+    @T.prim_func(s_tir=True)
+    def ramp_compare(B: T.Buffer((16,), "int32")):
+        for i in T.vectorized(16):
+            B[i] = T.Select(i * 3 + 5 < 29, i * 3 + 5, -1)
+
+    target = tvm.target.Target(
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        }
+    )
+
+    with target:
+        f = tvm.tirx.build(ramp_compare, target)
+
+    assembly = f.inspect_source("asm")
+    assert "vid.v" in assembly
+    assert re.search(r"\bvmul\.v", assembly)
+    assert re.search(r"\bvadd\.v", assembly)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_rocm.py

@@ -0,0 +1,221 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F841
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_inf_nan():
+    def check_inf_nan(n, value, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i_0 in T.thread_binding(1, thread="blockIdx.x"):
+                    for i_1 in T.thread_binding(128, thread="threadIdx.x"):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(1, i_0 * 128 + i_1)
+                            T.where(i_0 * 128 + i_1 < 1)
+                            T.reads()
+                            T.writes(C[v_i])
+                            C[v_i] = T.Cast(dtype, value)
+
+        fun = tvm.compile(Module, "rocm")
+
+        def run_and_check():
+            dev = tvm.rocm(0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_inf_nan(1, -float("inf"), "float32")
+    check_inf_nan(1, -float("inf"), "float64")
+    check_inf_nan(1, float("inf"), "float32")
+    check_inf_nan(1, float("inf"), "float64")
+    check_inf_nan(1, float("nan"), "float32")
+    check_inf_nan(1, float("nan"), "float64")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_copy():
+    def check_rocm(dtype, n):
+        def run_and_check():
+            dev = tvm.rocm(0)
+            a_np = np.random.uniform(size=(n,)).astype(dtype)
+            a = tvm.runtime.empty((n,), dtype, dev).copyfrom(a_np)
+            b_np = a.numpy()
+            tvm.testing.assert_allclose(a_np, b_np)
+            tvm.testing.assert_allclose(a_np, a.numpy())
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    for _ in range(100):
+        dtype = np.random.choice(["float32", "float16", "int8", "int32"])
+        logN = np.random.randint(1, 15)
+        peturb = np.random.uniform(low=0.5, high=1.5)
+        check_rocm(dtype, int(peturb * (2**logN)))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_vectorize_add():
+    def check_rocm(dtype, n, lanes):
+        vec_dtype = f"{dtype}x{lanes}"
+        num_blocks = n // 4
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((n,), vec_dtype), B: T.Buffer((n,), vec_dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
+                    for i_1 in T.thread_binding(4, thread="threadIdx.x"):
+                        with T.sblock("B"):
+                            v_i = T.axis.spatial(n, i_0 * 4 + i_1)
+                            T.reads(A[v_i])
+                            T.writes(B[v_i])
+                            B[v_i] = A[v_i] + T.Broadcast(T.Cast(dtype, 1), lanes)
+
+        fun = tvm.compile(Module, target="rocm")
+
+        def run_and_check():
+            dev = tvm.rocm(0)
+            a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np.random.uniform(size=(n, lanes)))
+            c = tvm.runtime.empty((n,), vec_dtype, dev)
+            fun(a, c)
+            tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_rocm("float32", 64, 2)
+    check_rocm("float16", 64, 2)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_warp_shuffle():
+    @T.prim_func(s_tir=True)
+    def func(
+        A_handle: T.handle,
+    ):
+        A = T.match_buffer(A_handle, (32,), dtype="float32")
+
+        for bx in T.thread_binding(1, thread="blockIdx.x"):
+            for tx in T.thread_binding(32, thread="threadIdx.x"):
+                with T.sblock("test"):
+                    A_local = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                    mask = T.sblock_alloc_buffer((1,), "uint32", scope="local")
+                    t0 = T.sblock_alloc_buffer((1,), "float32", scope="local")
+
+                    A_local[0] = A[tx]
+                    A_local[0] = T.tvm_warp_shuffle(mask[0], A_local[0], 0, 32, 32)
+                    A[tx] = A_local[0]
+
+    mod = tvm.compile(func, target="rocm")
+
+    def run_and_check():
+        dev = tvm.rocm(0)
+        a = tvm.runtime.tensor(np.random.uniform(size=(32,)).astype("float32"), dev)
+        mod(a)
+        tvm.testing.assert_allclose(a.numpy(), np.ones((32,)) * a.numpy()[0])
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_vectorized_exp():
+    @T.prim_func(s_tir=True)
+    def func(
+        A_handle: T.handle,
+        B_handle: T.handle,
+    ):
+        A = T.match_buffer(A_handle, (4,), dtype="float32")
+        B = T.match_buffer(B_handle, (4,), dtype="float32")
+
+        for bx in T.thread_binding(1, thread="blockIdx.x"):
+            for tx in T.thread_binding(1, thread="threadIdx.x"):
+                with T.sblock("test"):
+                    for i in T.vectorized(0, 4):
+                        B[i] = T.exp2(A[i])
+
+    mod = tvm.compile(func, target="rocm")
+
+    def run_and_check():
+        dev = tvm.rocm(0)
+        a = tvm.runtime.tensor(np.ones((4,)).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.zeros((4,)).astype("float32"), dev)
+        mod(a, b)
+        tvm.testing.assert_allclose(b.numpy(), np.exp2(a.numpy()))
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_export_load_with_fallback(monkeypatch, tmp_path):
+    """Force the codegen wrapper into the fallback branch, then export+load+run."""
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+    host_lib = tvm.compile(Module, target="rocm")
+    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
+
+    lib_path = str(tmp_path / "lib.so")
+    host_lib.export_library(lib_path)
+    reloaded = tvm.runtime.load_module(lib_path)
+
+    a_np = np.random.uniform(size=(n,)).astype("float32")
+    b_np = np.zeros((n,), dtype="float32")
+
+    def run_and_check():
+        dev = tvm.rocm(0)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        reloaded["main"](a, b)
+        np.testing.assert_allclose(b.numpy(), a_np + 1.0, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_static_init.py

@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import ctypes
+
+import numpy as np
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_static_init():
+    @tvm.register_global_func("test_static_callback")
+    def test_cb(sh, A):
+        assert isinstance(sh, ctypes.c_void_p)
+        return sh
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def ramp(A: T.handle):
+            T.func_attr({"global_symbol": "ramp"})
+            n = T.int64()
+            Ab = T.match_buffer(A, (n,), "int64")
+            T.call_packed(
+                "test_static_callback",
+                T.call_intrin("handle", "tirx.tvm_static_handle"),
+                Ab.data,
+            )
+
+    mod = Module
+    f = tvm.driver.build(mod, target="llvm")
+    a = tvm.runtime.tensor(np.zeros(10, dtype="int64"))
+    f(a)
+
+
+if __name__ == "__main__":
+    test_static_init()

tests/python/codegen/test_target_codegen_vulkan.py

@@ -0,0 +1,699 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501, F841
+
+import re
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import ir as I_builder
+from tvm.script.ir_builder import tirx as T_builder
+from tvm.testing import env
+
+dtype = tvm.testing.parameter("float32", "int32", "float16", "int8")
+fuzz_seed = tvm.testing.parameter(range(25))
+
+
+# Explicitly specify a target, as this test is looking at the
+# generated shader code, and is not running on an actual device.
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled(
+        {
+            "kind": "vulkan",
+            "supports_int8": 1,
+            "supports_8bit_buffer": 1,
+            "supports_storage_buffer_storage_class": 1,
+            "supports_float16": 1,
+            "supports_16bit_buffer": 1,
+        }
+    ),
+    reason="vulkan not enabled",
+)
+def test_vector_comparison(dtype):
+    target = {
+        "kind": "vulkan",
+        "supports_int8": 1,
+        "supports_8bit_buffer": 1,
+        "supports_storage_buffer_storage_class": 1,
+        "supports_float16": 1,
+        "supports_16bit_buffer": 1,
+    }
+    target = tvm.target.Target(target)
+    zero = tvm.tirx.const(0, dtype)
+    one = tvm.tirx.const(1, dtype)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1024,), dtype), B: T.Buffer((1024,), dtype)):
+            for i_0 in T.thread_binding(8, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    for i_2 in T.vectorized(4):
+                        with T.sblock("B"):
+                            v_i = T.axis.spatial(1024, i_0 * 128 + i_1 * 4 + i_2)
+                            B[v_i] = T.Select(A[v_i] >= zero, A[v_i] + one, zero)
+
+    # Build
+    f = tvm.tirx.build(Module, target=target)
+
+    # Verify we generate the boolx4 type declaration and the OpSelect
+    # v4{float,half,int} instruction
+    assembly = f.imports[0].inspect_source()
+    matches = re.findall("%v4bool = OpTypeVector %bool 4", assembly)
+    assert len(matches) == 1
+    matches = re.findall("OpSelect %v4.*", assembly)
+    assert len(matches) == 1
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        "llvm",
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("rocm", marks=pytest.mark.gpu),
+        pytest.param("vulkan", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+        pytest.param("opencl", marks=pytest.mark.gpu),
+        pytest.param("nvptx", marks=pytest.mark.gpu),
+    ],
+)
+def test_array_copy(target, dtype, fuzz_seed):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    np.random.seed(fuzz_seed)
+
+    log_arr_size = np.random.uniform(low=np.log(1), high=np.log(32768))
+    arr_size = np.exp(log_arr_size).astype(int)
+    a_np = np.random.uniform(size=(arr_size,)).astype(dtype)
+
+    def run_and_check():
+        dev = tvm.device(target)
+        a = tvm.runtime.empty((arr_size,), dtype, dev).copyfrom(a_np)
+        tvm.testing.assert_allclose(a_np, a.numpy())
+
+    if target == "llvm":
+        run_and_check()
+    else:
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_array_vectorize_add(dtype):
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    arr_size = 64
+    lanes = 2
+
+    vec_dtype = f"{dtype}x{lanes}"
+    one = tvm.tirx.const(1, vec_dtype)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((64,), vec_dtype), B: T.Buffer((64,), vec_dtype)):
+            for i_0 in T.thread_binding(16, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(4, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(64, i_0 * 4 + i_1)
+                        B[v_i] = A[v_i] + one
+
+    f = tvm.compile(Module, target=target)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        a = tvm.runtime.empty((arr_size,), vec_dtype, dev).copyfrom(
+            np.random.uniform(size=(arr_size, lanes))
+        )
+        c = tvm.runtime.empty((arr_size,), vec_dtype, dev)
+        f(a, c)
+        tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vulkan_bool_load():
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    arr_size = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1024,), "bool"), B: T.Buffer((1024,), "int32")):
+            for i_0 in T.thread_binding(8, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(128, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(1024, i_0 * 128 + i_1)
+                        B[v_i] = T.Cast("int32", A[v_i])
+
+    f = tvm.compile(Module, target=target)
+
+    a_np = np.random.uniform(size=arr_size) > 0.5
+    b_np = np.zeros((arr_size,), dtype="int32")
+    ref = a_np.astype(np.int32)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        f(a, b)
+        tvm.testing.assert_allclose(b.numpy(), ref)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+vulkan_parameter_impl = tvm.testing.parameter("push_constants", "ubo")
+vulkan_parameter_dtype = tvm.testing.parameter("int32", "float32", "int64")
+
+
+# Only run on vulkan because extremely large numbers of input
+# parameters can crash cuda/llvm compiler.
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vulkan_constant_passing(vulkan_parameter_impl, vulkan_parameter_dtype):
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    dtype = vulkan_parameter_dtype
+
+    if not target.attrs.get("supports_int64", False):
+        pytest.xfail("Vulkan target does not support Int64 variables")
+
+    # f_add has 3+num_int_params scalar parameters.  The other three
+    # are length_n, stride1, and stride2.
+    if vulkan_parameter_impl == "push_constants":
+        # 4 params, 32 bytes.  Within 128-byte spec-guaranteed size of
+        # push constants.  Uses push constants.
+        num_int_params = 1
+    else:
+        # 24 params, 192 bytes.  May be above spec-guaranteed size of 128
+        # bytes for push constants.  Uses either push constants or UBO,
+        # depending on the device.
+        max_push_constants_size = int(target.attrs.get("max_push_constants_size", 128))
+        max_int_params_in_push = max_push_constants_size // 8 - 3
+        num_int_params = max_int_params_in_push + 1
+
+    # Build IRModule programmatically since num_int_params is dynamic
+    with IRBuilder() as ib:
+        with I_builder.ir_module():
+            with T_builder.prim_func():
+                T_builder.func_name("main")
+                scalar_vars = []
+                for i in range(num_int_params):
+                    v = T_builder.arg(f"scale{i}", tvm.tirx.Var("", dtype))
+                    scalar_vars.append(v)
+                var_A = T_builder.arg("var_A", T_builder.handle())
+                var_B = T_builder.arg("var_B", T_builder.handle())
+                T_builder.func_attr({"tirx.noalias": True})
+                n_var = T_builder.int32()
+                A = T_builder.match_buffer(var_A, (n_var,), dtype)
+                B = T_builder.match_buffer(var_B, (n_var,), dtype)
+                scalar_sum = scalar_vars[0]
+                for s in scalar_vars[1:]:
+                    scalar_sum = scalar_sum + s
+                with T_builder.thread_binding(
+                    tvm.tirx.ceildiv(n_var, 64), thread="blockIdx.x"
+                ) as i_0:
+                    with T_builder.thread_binding(64, thread="threadIdx.x") as i_1:
+                        with T_builder.sblock("B"):
+                            v_i = T_builder.axis.spatial(n_var, i_0 * 64 + i_1)
+                            T_builder.where(i_0 * 64 + i_1 < n_var)
+                            T_builder.reads(A[v_i])
+                            T_builder.writes(B[v_i])
+                            T_builder.buffer_store(B, scalar_sum + A[v_i], [v_i])
+    mod = ib.get()
+    f_add = tvm.compile(mod, target=target)
+
+    n = 1024
+    scalars = np.array([1 for _ in range(num_int_params)]).astype(dtype)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        a = tvm.runtime.tensor(np.random.uniform(size=n).astype(dtype), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
+        f_add(*scalars, a, b)
+        tvm.testing.assert_allclose(a.numpy() + sum(scalars), b.numpy())
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vulkan_while_if():
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    n = 1
+    dtype = "int32"
+
+    @T.prim_func(s_tir=True)
+    def while_if_gpu(A: T.Buffer((1,), "int32"), B: T.Buffer((1,), "int32")):
+        for bx in T.thread_binding(1, thread="blockIdx.x"):
+            iterations = T.decl_buffer((1,), "int32", scope="local")
+            iterations[0] = 0
+            B[0] = 0
+            while iterations[0] < T.if_then_else(A[0] > 0, 10, 20):
+                iterations[0] = iterations[0] + 1
+                B[0] = B[0] + iterations[0]
+
+    mod = tvm.IRModule.from_expr(while_if_gpu.with_attr("target", target))
+    compiled_func = tvm.compile(mod, target=target)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        for input_value, expected in [(5, [55]), (-5, [210])]:
+            a = tvm.runtime.tensor(np.array([input_value], dtype=dtype), dev)
+            b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
+            compiled_func(a, b)
+            tvm.testing.assert_allclose(b.numpy(), expected)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vulkan_local_threadidx():
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    n = 32
+
+    @T.prim_func(s_tir=True)
+    def local_threadidx_func(A: T.Buffer((32,), "int32"), B: T.Buffer((32,), "int32")):
+        # First block with thread extent 16
+        for _ in range(1):
+            for tx in T.thread_binding(16, thread="threadIdx.x"):
+                B[tx + 0] = A[tx + 0]
+        # Second block with thread extent 16
+        for _ in range(1):
+            for tx in T.thread_binding(16, thread="threadIdx.x"):
+                B[tx + 16] = A[tx + 16]
+
+    mod = tvm.IRModule.from_expr(local_threadidx_func)
+    func = tvm.compile(mod, target=target)
+
+    a_np = np.arange(n).astype(dtype="int32")
+    b_np = np.zeros((n,), dtype="int32")
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        func(a, b)
+        tvm.testing.assert_allclose(b.numpy(), a_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vectorized_index_ramp():
+    """Test vectorized copy with ramp indices (load N values, write to N locations)"""
+    target = {"kind": "vulkan", "from_device": 0}
+    n = 4
+    ramp_index = tvm.tirx.Ramp(0, 1, 4)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            A = T.match_buffer(var_A, (n,), "int32", offset_factor=1)
+            B = T.match_buffer(var_B, (n,), "int32", offset_factor=1)
+            with T.sblock("compute"):
+                T.reads()
+                T.writes()
+                bx = T.launch_thread("blockIdx.x", 1)
+                B[ramp_index] = A[ramp_index]
+
+    f = tvm.compile(Module, target=target)
+
+    a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
+    b_np = np.zeros(n, dtype="int32")
+
+    def run_and_check():
+        dev = tvm.device(target["kind"])
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        f(a, b)
+        tvm.testing.assert_allclose(b.numpy(), a_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vectorized_index_broadcast():
+    """Test broadcast index (load 1 value, write to N locations)"""
+    target = {"kind": "vulkan", "from_device": 0}
+    n = 4
+    broadcast_index = tvm.tirx.Broadcast(0, 4)
+    ramp_index = tvm.tirx.Ramp(0, 1, 4)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            A = T.match_buffer(var_A, (n,), "int32", offset_factor=1)
+            B = T.match_buffer(var_B, (n,), "int32", offset_factor=1)
+            with T.sblock("compute"):
+                T.reads()
+                T.writes()
+                bx = T.launch_thread("blockIdx.x", 1)
+                # Load from broadcast index (single element), store to ramp index
+                B[ramp_index] = A[broadcast_index]
+
+    f = tvm.compile(Module, target=target)
+
+    a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
+    b_np = np.zeros(n, dtype="int32")
+
+    def run_and_check():
+        dev = tvm.device(target["kind"])
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        f(a, b)
+        # All elements of b should be a[0] (broadcast load)
+        tvm.testing.assert_allclose(b.numpy(), np.full(n, a_np[0]))
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_negative_operand_divmod():
+    """Test handling of negative offsets to floormod/floordiv
+
+    Even though the SPIR-V spec states that OpSRem and OpSMod can give
+    the signed modulo, the Vulkan spec states that any use of negative
+    operands is undefined behavior.  This test starts with negative
+    operands to floordiv, validating that they are simplified into the
+    corresponding positive operands, such that the final TIR can be
+    expressed using only positive operands.
+
+    SPIR-V: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpSRem
+    Vulkan: https://registry.khronos.org/vulkan/specs/1.3/html/chap37.html#spirvenv-op-prec
+    """
+    target = {"kind": "vulkan", "from_device": 0}
+
+    N = 32
+    offset = 16
+    divisor = 5
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((N, 2), "int32")):
+        for i in T.thread_binding(N, thread="threadIdx.x"):
+            with T.sblock("A"):
+                v_i = T.axis.spatial(N, i)
+                A[v_i, 0] = T.floordiv(v_i - offset, divisor)
+                A[v_i, 1] = T.floormod(v_i - offset, divisor)
+
+    built = tvm.compile(func, target=target)
+
+    def run_and_check():
+        dev = tvm.device(target["kind"])
+        a_dev = tvm.runtime.empty([N, 2], "int32", dev)
+        built(a_dev)
+        a = a_dev.numpy()
+        np.testing.assert_array_equal(a[:, 0], (np.arange(N) - offset) // divisor)
+        np.testing.assert_array_equal(a[:, 1], (np.arange(N) - offset) % divisor)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("out_dtype", ["float32", "float16"])
+def test_cooperative_matrix(out_dtype):
+    M, N, K = 16, 16, 32
+
+    # fmt: off
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(X: T.Buffer((16, 32), "float16"), W: T.Buffer((32, 16), "float16"), compute: T.Buffer((16, 16), out_dtype)):
+            T.func_attr({"tirx.noalias": True})
+            X_shared = T.sblock_alloc_buffer((16, 32), "float16", scope="shared")
+            W_shared = T.sblock_alloc_buffer((32, 16), "float16", scope="shared")
+            X_shared_wmma_matrix_a = T.sblock_alloc_buffer((16, 32), "float16", scope="wmma.matrix_a")
+            W_shared_wmma_matrix_b = T.sblock_alloc_buffer((32, 16), "float16", scope="wmma.matrix_b")
+            compute_wmma_accumulator = T.sblock_alloc_buffer((16, 16), out_dtype, scope="wmma.accumulator")
+            for i_0_j_0_fused in T.thread_binding(1, thread="blockIdx.x"):
+                with T.sblock("compute_init_o"):
+                    v_i_o = T.axis.spatial(1, 0)
+                    v_j_o = T.axis.spatial(1, 0)
+                    T.reads()
+                    T.writes(compute_wmma_accumulator[0:16, 0:16])
+                    C = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
+                    T.tvm_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0.0))
+                for k_0 in range(2):
+                    for ax0_ax1_fused_0 in range(2):
+                        for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
+                            for ax0_ax1_fused_2 in T.vectorized(4):
+                                with T.sblock("X_shared"):
+                                    v0 = T.axis.spatial(16, (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 16)
+                                    v1 = T.axis.spatial(32, k_0 * 16 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 16)
+                                    T.reads(X[v0, v1])
+                                    T.writes(X_shared[v0, v1])
+                                    X_shared[v0, v1] = X[v0, v1]
+                    for ax0_ax1_fused_0 in range(2):
+                        for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
+                            for ax0_ax1_fused_2 in T.vectorized(4):
+                                with T.sblock("W_shared"):
+                                    v0 = T.axis.spatial(32, k_0 * 16 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 16)
+                                    v1 = T.axis.spatial(16, (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 16)
+                                    T.reads(W[v0, v1])
+                                    T.writes(W_shared[v0, v1])
+                                    W_shared[v0, v1] = W[v0, v1]
+                    for ax0_0 in T.unroll(1):
+                        for ax1_0 in T.unroll(1):
+                            with T.sblock("X_shared_wmma.matrix_a_o"):
+                                v0_o = T.axis.spatial(1, ax0_0)
+                                v1_o = T.axis.spatial(2, k_0 + ax1_0)
+                                T.reads(X_shared[0:16, v1_o * 16:v1_o * 16 + 16])
+                                T.writes(X_shared_wmma_matrix_a[0:16, v1_o * 16:v1_o * 16 + 16])
+                                A = T.match_buffer(X_shared[0:16, v1_o * 16:v1_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=16)
+                                C = T.match_buffer(X_shared_wmma_matrix_a[0:16, v1_o * 16:v1_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_a", offset_factor=16)
+                                T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
+                    for ax0_0 in T.unroll(1):
+                        for ax1_0 in T.unroll(1):
+                            with T.sblock("W_shared_wmma.matrix_b_o"):
+                                v0_o = T.axis.spatial(2, k_0 + ax0_0)
+                                v1_o = T.axis.spatial(1, ax1_0)
+                                T.reads(W_shared[v0_o * 16:v0_o * 16 + 16, 0:16])
+                                T.writes(W_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, 0:16])
+                                A = T.match_buffer(W_shared[v0_o * 16:v0_o * 16 + 16, 0:16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=16)
+                                C = T.match_buffer(W_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, 0:16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_b", offset_factor=16)
+                                T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
+                    with T.sblock("compute_update_o"):
+                        v_i_o = T.axis.spatial(1, 0)
+                        v_j_o = T.axis.spatial(1, 0)
+                        v_k_o = T.axis.reduce(2, k_0)
+                        T.reads(compute_wmma_accumulator[0:16, 0:16], X_shared_wmma_matrix_a[0:16, v_k_o * 16:v_k_o * 16 + 16], W_shared_wmma_matrix_b[v_k_o * 16:v_k_o * 16 + 16, 0:16])
+                        T.writes(compute_wmma_accumulator[0:16, 0:16])
+                        A = T.match_buffer(X_shared_wmma_matrix_a[0:16, v_k_o * 16:v_k_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="wmma.matrix_a", offset_factor=16)
+                        B = T.match_buffer(W_shared_wmma_matrix_b[v_k_o * 16:v_k_o * 16 + 16, 0:16], (16, 16), "float16", strides=("B_s0", "B_s1"), scope="wmma.matrix_b", offset_factor=16)
+                        C = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
+                        T.tvm_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
+                with T.sblock("compute_wmma.accumulator_o"):
+                    v0_o = T.axis.spatial(1, 0)
+                    v1_o = T.axis.spatial(1, 0)
+                    T.reads(compute_wmma_accumulator[0:16, 0:16])
+                    T.writes(compute[0:16, 0:16])
+                    A = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("A_s0", "A_s1"), scope="wmma.accumulator", offset_factor=16)
+                    C = T.match_buffer(compute[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), offset_factor=16)
+                    T.tvm_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(out_dtype), C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
+    # fmt: on
+
+    target = {"kind": "vulkan", "from_device": 0}
+    tgt_attrs = tvm.target.Target(target).attrs
+
+    if tgt_attrs.get("supports_cooperative_matrix"):
+        f = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device("vulkan", 0)
+            A = tvm.runtime.tensor(np.random.randn(M, K).astype("float16"), dev)
+            B = tvm.runtime.tensor(np.random.randn(K, N).astype("float16"), dev)
+            C = tvm.runtime.tensor(np.random.randn(M, N).astype(out_dtype), dev)
+            f(A, B, C)
+            A_np = A.numpy()
+            B_np = B.numpy()
+            ref = np.dot(A_np.astype("float32"), B_np.astype("float32"))
+            tvm.testing.assert_allclose(C.numpy(), ref, rtol=1e-2, atol=1e-2)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
+def test_codegen_decl_buffer():
+    """The codegen should accept DeclBuffer nodes in its input"""
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def kernel():
+            T.func_attr({"calling_conv": 2, "global_symbol": "kernel", "tirx.noalias": True})
+            A = T.alloc_buffer((256,), dtype="float32", scope="local")
+            A_buf = T.decl_buffer([256], dtype="float32", scope="local", data=A.data)
+
+    target = tvm.target.Target("vulkan")
+    vulkan_codegen = tvm.get_global_func("target.build.vulkan")
+    vulkan_codegen(Module, target)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
+def test_codegen_static_shared_memory():
+    """The codegen should accept static shared/workgroup allocations."""
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+            A_shared = T.alloc_buffer((128,), dtype="float32", scope="shared")
+
+            for bx in T.thread_binding(1, thread="blockIdx.x"):
+                for tx in T.thread_binding(128, thread="threadIdx.x"):
+                    A_shared[tx] = A[tx]
+                    B[tx] = A_shared[tx]
+
+    tvm.compile(Module, target="vulkan")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
+def test_unary():
+    test_funcs = [
+        (tvm.tirx.sin, lambda x: np.sin(x)),
+        (tvm.tirx.cos, lambda x: np.cos(x)),
+        (tvm.tirx.tan, lambda x: np.tan(x)),
+        (tvm.tirx.sinh, lambda x: np.sinh(x)),
+        (tvm.tirx.cosh, lambda x: np.cosh(x)),
+        (tvm.tirx.tanh, lambda x: np.tanh(x)),
+        (tvm.tirx.asin, lambda x: np.arcsin(x)),
+        (tvm.tirx.acos, lambda x: np.arccos(x)),
+        (tvm.tirx.atan, lambda x: np.arctan(x)),
+        (tvm.tirx.asinh, lambda x: np.arcsinh(x)),
+        (tvm.tirx.acosh, lambda x: np.arccosh(x)),
+        (tvm.tirx.atanh, lambda x: np.arctanh(x)),
+    ]
+
+    def run_test(tvm_intrin, np_func):
+        n = 16
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(var_A: T.handle, var_B: T.handle):
+                m = T.int32()
+                A = T.match_buffer(var_A, (m,), "float32")
+                B = T.match_buffer(var_B, (m,), "float32")
+                for i_0 in T.thread_binding((m + 63) // 64, thread="blockIdx.x"):
+                    for i_1 in T.thread_binding(64, thread="threadIdx.x"):
+                        with T.sblock("B"):
+                            v_i = T.axis.spatial(m, i_0 * 64 + i_1)
+                            T.where(i_0 * 64 + i_1 < m)
+                            T.reads(A[v_i])
+                            T.writes(B[v_i])
+                            B[v_i] = tvm_intrin(A[v_i])
+
+        target = tvm.target.Target("vulkan")
+        func = tvm.compile(Module, target=target)
+
+        if tvm_intrin in [tvm.tirx.asin, tvm.tirx.acos]:
+            data = np.random.uniform(-1.0, 1.0, size=n)
+        elif tvm_intrin == tvm.tirx.atanh:
+            data = np.random.uniform(-0.999, 0.999, size=n)
+        elif tvm_intrin == tvm.tirx.acosh:
+            data = np.random.uniform(1.0, 5.0, size=n)
+        else:
+            data = np.random.uniform(0.1, 0.9, size=n)
+
+        def run_and_check():
+            dev = tvm.device(target.kind.name, 0)
+            a = tvm.runtime.tensor(data.astype("float32"), dev)
+            b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+            func(a, b)
+            tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=1e-3, rtol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    for func in test_funcs:
+        run_test(*func)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
+def test_export_load_with_fallback(monkeypatch, tmp_path):
+    """Force the codegen wrapper into the fallback branch, then export."""
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+    host_lib = tvm.compile(Module, target="vulkan")
+    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
+
+    lib_path = str(tmp_path / "lib.so")
+    host_lib.export_library(lib_path)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_x86.py

@@ -0,0 +1,88 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E741
+import platform
+import re
+
+import pytest
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+llvm_version = tvm.target.codegen.llvm_version_major()
+machine = platform.machine()
+
+if machine not in ["x86_64", "AMD64", "amd64"]:
+    pytest.skip(f"Requires x86_64, but machine is {machine}", allow_module_level=True)
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+@pytest.mark.skipif(llvm_version < 6, reason=f"Requires LLVM 6+, got {llvm_version}")
+def test_fp16_to_fp32():
+    def fp16_to_fp32(target, width, match=None, not_match=None):
+        elements = 64
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(
+                A: T.Buffer((elements, width), "float16"),
+                B: T.Buffer((elements, width), "float32"),
+            ):
+                T.func_attr({"tirx.noalias": True})
+                for i0 in range(elements):
+                    for i1 in T.vectorized(width):
+                        with T.sblock("B"):
+                            v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+                            T.reads(A[v_i0, v_i1])
+                            T.writes(B[v_i0, v_i1])
+                            B[v_i0, v_i1] = T.Cast("float32", A[v_i0, v_i1])
+
+        f = tvm.tirx.build(Module, target=target)
+
+        assembly = f.inspect_source("asm").splitlines()
+        if match:
+            matches = [l for l in assembly if re.search(match, l)]
+            assert matches
+        if not_match:
+            not_matches = [l for l in assembly if re.search(not_match, l)]
+            assert not not_matches
+
+    fp16_to_fp32({"kind": "llvm", "mcpu": "skylake-avx512"}, 15, match="vcvtph2ps.*mm")
+    fp16_to_fp32({"kind": "llvm", "mcpu": "skylake-avx512"}, 16, match="vcvtph2ps.*mm")
+    fp16_to_fp32({"kind": "llvm", "mcpu": "skylake-avx512"}, 17, match="vcvtph2ps.*mm")
+    fp16_to_fp32({"kind": "llvm", "mcpu": "skylake-avx512"}, 49, match="vcvtph2ps.*mm")
+    fp16_to_fp32(
+        {"kind": "llvm", "mcpu": "skylake-avx512", "mattr": ["-avx512f"]}, 49, match="vcvtph2ps.*mm"
+    )
+    fp16_to_fp32(
+        {"kind": "llvm", "mcpu": "skylake-avx512", "mattr": ["-f16c", "-avx512f"]},
+        49,
+        not_match="vcvtph2ps",
+    )
+    fp16_to_fp32({"kind": "llvm", "mcpu": "core-avx2"}, 8, match="vcvtph2ps.*mm")
+    fp16_to_fp32({"kind": "llvm", "mcpu": "core-avx2"}, 9, match="vcvtph2ps.*mm")
+    fp16_to_fp32("llvm", 9, not_match="vcvtph2ps")
+
+
+is_32bit = platform.architecture()[0] == "32bit"
+
+
+if __name__ == "__main__":
+    test_fp16_to_fp32()

tests/python/conftest.py

@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Configure pytest for TVM's Python test suite."""
+
+import os
+from pathlib import Path
+
+
+def pytest_sessionstart():
+    if os.getenv("CI", "") == "true":
+        from tvm.testing.utils import (
+            install_request_hook,  # pylint: disable=import-outside-toplevel
+        )
+
+        install_request_hook(Path(__file__).with_name("request_hook.py"))

tests/python/contrib/pickle_memoize_script.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import sys
+
+import tvm
+
+
+@tvm.contrib.pickle_memoize.memoize("test_memoize_save_data", save_at_exit=True)
+def get_data_saved():
+    return 42
+
+
+@tvm.contrib.pickle_memoize.memoize("test_memoize_transient_data", save_at_exit=False)
+def get_data_transient():
+    return 42
+
+
+def main():
+    assert len(sys.argv) == 3, "Expect arguments SCRIPT NUM_SAVED NUM_TRANSIENT"
+
+    num_iter_saved = int(sys.argv[1])
+    num_iter_transient = int(sys.argv[2])
+
+    for _ in range(num_iter_saved):
+        get_data_saved()
+    for _ in range(num_iter_transient):
+        get_data_transient()
+
+
+if __name__ == "__main__":
+    main()

tests/python/contrib/test_android/__init__.py

@@ -0,0 +1,19 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Testing infrastructure for Android"""

tests/python/contrib/test_android/infrastructure.py

@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+
+"""Android testing infrastructure"""
+
+import os
+
+import tvm
+from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig, RPCRunner
+
+
+def get_rpc_runner() -> tvm.s_tir.meta_schedule.runner.RPCRunner:
+    if (
+        "TVM_TRACKER_HOST" in os.environ
+        and "TVM_TRACKER_PORT" in os.environ
+        and "RPC_DEVICE_KEY" in os.environ
+    ):
+        rpc_host = os.environ["TVM_TRACKER_HOST"]
+        rpc_port = int(os.environ["TVM_TRACKER_PORT"])
+        rpc_key = os.environ["RPC_DEVICE_KEY"]
+    else:
+        raise Exception("Please initialize environment variables for using RPC tracker")
+
+    rpc_config = RPCConfig(
+        tracker_host=rpc_host,
+        tracker_port=rpc_port,
+        tracker_key=rpc_key,
+        session_priority=1,
+        session_timeout_sec=100,
+    )
+    evaluator_config = EvaluatorConfig(
+        number=1,
+        repeat=1,
+        min_repeat_ms=0,
+    )
+    return RPCRunner(rpc_config, evaluator_config)
+
+
+def get_android_gpu_target() -> tvm.target.Target:
+    """Creates a Android GPU target"""
+    target_c = "opencl"
+    target_h = {"kind": "llvm", "mtriple": "arm64-linux-android"}
+    return tvm.target.Target(target_c, host=target_h)

tests/python/contrib/test_android/test_meta_schedule.py

@@ -0,0 +1,78 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+
+"""Test rpc based launcher for Android"""
+
+import tempfile
+
+import numpy as np
+import pytest
+
+pytest.importorskip("scipy")  # tvm.topi.testing imports scipy
+
+import tvm.testing
+import tvm.topi.testing
+from tvm.s_tir import meta_schedule as ms
+from tvm.s_tir.meta_schedule.builder import LocalBuilder
+from tvm.script import tirx as T
+
+from .infrastructure import get_android_gpu_target, get_rpc_runner
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@pytest.mark.skip("Integration test")
+def test_tune_tir_on_android():
+    """Test tune_tir on Android through RPC."""
+    max_workers = 4
+    builder = LocalBuilder(
+        f_export="s_tir.meta_schedule.builder.export_ndk", max_workers=max_workers
+    )
+    runner = get_rpc_runner()
+    target = get_android_gpu_target()
+    with tempfile.TemporaryDirectory() as work_dir:
+        database = ms.tir_integration.tune_tir(
+            mod=matmul,
+            target=target,
+            work_dir=work_dir,
+            max_trials_global=32,
+            num_trials_per_iter=16,
+            builder=builder,
+            runner=runner,
+        )
+        sch = ms.tir_integration.compile_tir(database, matmul, target)
+        if sch is None:
+            print("No valid schedule found!")
+        else:
+            sch.mod.show()
+            sch.trace.show()
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/contrib/test_cblas.py

@@ -0,0 +1,301 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Configure pytest"""
+
+import numpy as np
+import pytest
+
+pytest.importorskip("scipy")  # tvm.topi.testing imports scipy
+
+import tvm
+import tvm.testing
+import tvm.topi.testing
+from tvm import te
+from tvm.contrib import cblas, dnnl, mkl
+
+
+def verify_matmul_add(
+    matrix_m, matrix_l, matrix_n, lib, transa=False, transb=False, dtype="float32"
+):
+    """Tests matmul+add op"""
+    bias = te.var("bias", dtype=dtype)
+    ashape = (matrix_l, matrix_n) if transa else (matrix_n, matrix_l)
+    bshape = (matrix_m, matrix_l) if transb else (matrix_l, matrix_m)
+    input1_data = te.placeholder(ashape, name="input1_data", dtype=dtype)
+    input2_data = te.placeholder(bshape, name="input2_data", dtype=dtype)
+    matmul_result = lib.matmul(input1_data, input2_data, transa, transb)
+    final_result = te.compute(
+        matmul_result.shape, lambda i, j: matmul_result[i, j] + bias, name="final_result"
+    )
+
+    def get_numpy(a, b, matrix_bias, transa, transb):
+        if transa:
+            a = a.transpose()
+        if transb:
+            b = b.transpose()
+        return np.dot(a, b) + matrix_bias
+
+    def compiling(f, name="test_matmul_add", ext=".so"):
+        path = name + ext
+        f.export_library(path)
+        mod = tvm.runtime.load_module(path)
+        f = mod[name]
+        return f
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func(lib.__name__ + ".matmul", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        name = "test_matmul_add"
+        f = tvm.compile(
+            te.create_prim_func([input1_data, input2_data, final_result, bias]).with_attr(
+                "global_symbol", name
+            ),
+            target=target,
+        )
+        if target == "c":
+            f = compiling(f, name)
+        matrix_input1 = tvm.runtime.tensor(
+            np.random.uniform(size=ashape).astype(input1_data.dtype), dev
+        )
+        matrix_input2 = tvm.runtime.tensor(
+            np.random.uniform(size=bshape).astype(input2_data.dtype), dev
+        )
+        matrix_result = tvm.runtime.tensor(
+            np.zeros((matrix_n, matrix_m), dtype=final_result.dtype), dev
+        )
+        matrix_bias = 10.0
+        f(matrix_input1, matrix_input2, matrix_result, matrix_bias)
+        tvm.testing.assert_allclose(
+            matrix_result.numpy(),
+            get_numpy(matrix_input1.numpy(), matrix_input2.numpy(), matrix_bias, transa, transb),
+            rtol=1e-5,
+        )
+
+    verify("llvm")
+    verify("c")
+
+
+def test_matmul_add():
+    """Tests of matmul+add op"""
+    verify_matmul_add(235, 128, 1024, cblas)
+    verify_matmul_add(235, 128, 1024, cblas, True, False)
+    verify_matmul_add(235, 128, 1024, cblas, False, True)
+    verify_matmul_add(235, 128, 1024, cblas, True, True)
+    verify_matmul_add(235, 128, 1024, mkl)
+    verify_matmul_add(235, 128, 1024, mkl, True, False)
+    verify_matmul_add(235, 128, 1024, mkl, False, True)
+    verify_matmul_add(235, 128, 1024, mkl, True, True)
+    verify_matmul_add(235, 128, 1024, dnnl)
+    verify_matmul_add(235, 128, 1024, dnnl, True, False)
+    verify_matmul_add(235, 128, 1024, dnnl, False, True)
+    verify_matmul_add(235, 128, 1024, dnnl, True, True)
+    verify_matmul_add(1, 16, 4, cblas)
+    verify_matmul_add(1, 16, 3, cblas, True, False)
+    verify_matmul_add(1, 16, 3, cblas, False, False)
+    verify_matmul_add(1, 16, 3, cblas, True, True)
+    verify_matmul_add(1, 16, 4, mkl)
+    verify_matmul_add(1, 16, 3, mkl, True, False)
+    verify_matmul_add(1, 16, 3, mkl, False, False)
+    verify_matmul_add(1, 16, 3, mkl, True, True)
+    verify_matmul_add(1, 16, 4, dnnl)
+    verify_matmul_add(1, 16, 3, dnnl, True, False)
+    verify_matmul_add(1, 16, 3, dnnl, False, False)
+    verify_matmul_add(1, 16, 3, dnnl, True, True)
+
+
+def verify_quantized_matmul_add(matrix_m, matrix_l, matrix_n, transa=False, transb=False):
+    """Tests quantized matmul+add op"""
+    if not tvm.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True):
+        pytest.skip("Quantized dense is supported only for MKL. TVM GPU CI uses openblas")
+    data_dtype = "uint8"
+    kernel_dtype = "int8"
+    out_dtype = "int32"
+    bias = te.var("bias", dtype=out_dtype)
+    ashape = (matrix_l, matrix_n) if transa else (matrix_n, matrix_l)
+    bshape = (matrix_m, matrix_l) if transb else (matrix_l, matrix_m)
+    input1_data = te.placeholder(ashape, name="input1_data", dtype=data_dtype)
+    input2_data = te.placeholder(bshape, name="input2_data", dtype=kernel_dtype)
+    matmul_result = mkl.matmul_u8s8s32(input1_data, input2_data, transa, transb, dtype=out_dtype)
+    final_result = te.compute(
+        matmul_result.shape, lambda i, j: matmul_result[i, j] + bias, name="final_result"
+    )
+
+    def get_numpy(a, b, matrix_bias, transa, transb):
+        if transa:
+            a = a.transpose()
+        if transb:
+            b = b.transpose()
+        return np.dot(a, b) + matrix_bias
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        f = tvm.compile(
+            te.create_prim_func([input1_data, input2_data, final_result, bias]), target=target
+        )
+        matrix_input1 = tvm.runtime.tensor(
+            np.random.randint(low=0, high=50, size=ashape).astype(input1_data.dtype), dev
+        )
+        matrix_input2 = tvm.runtime.tensor(
+            np.random.randint(low=0, high=50, size=bshape).astype(input2_data.dtype), dev
+        )
+        matrix_result = tvm.runtime.tensor(
+            np.zeros((matrix_n, matrix_m), dtype=final_result.dtype), dev
+        )
+        matrix_bias = 10
+        f(matrix_input1, matrix_input2, matrix_result, matrix_bias)
+        tvm.testing.assert_allclose(
+            matrix_result.numpy(),
+            get_numpy(
+                matrix_input1.numpy().astype("int32"),
+                matrix_input2.numpy().astype("int32"),
+                matrix_bias,
+                transa,
+                transb,
+            ),
+            rtol=1e-5,
+        )
+
+    verify()
+
+
+def test_quantized_matmul_add():
+    """Tests of quantized matmul+add op"""
+    verify_quantized_matmul_add(235, 128, 1024)
+    verify_quantized_matmul_add(235, 128, 1024, True, False)
+    verify_quantized_matmul_add(235, 128, 1024, False, True)
+    verify_quantized_matmul_add(235, 128, 1024, True, True)
+    verify_quantized_matmul_add(1, 16, 4)
+    verify_quantized_matmul_add(1, 16, 3, True, False)
+    verify_quantized_matmul_add(1, 16, 3, False, True)
+    verify_quantized_matmul_add(1, 16, 3, True, True)
+
+
+def verify_batch_matmul(
+    batch_a,
+    batch_b,
+    matrix_m,
+    matrix_l,
+    matrix_n,
+    lib,
+    transa=False,
+    transb=False,
+    dtype="float32",
+):
+    """Tests matmul op where matrices are in batch"""
+    batch = max(batch_a, batch_b)
+    ashape = (batch_a, matrix_l, matrix_n) if transa else (batch_a, matrix_n, matrix_l)
+    bshape = (batch_b, matrix_m, matrix_l) if transb else (batch_b, matrix_l, matrix_m)
+    input1_data = te.placeholder(ashape, name="input1_data", dtype=dtype)
+    input2_data = te.placeholder(bshape, name="input2_data", dtype=dtype)
+    matmul_result = lib.batch_matmul(input1_data, input2_data, transa, transb)
+    final_result = te.compute(
+        matmul_result.shape, lambda k, i, j: matmul_result[k, i, j], name="final_result"
+    )
+
+    def get_numpy(a, b, transa, transb):
+        if transa:
+            a = a.transpose(0, 2, 1)
+        if not transb:
+            b = b.transpose(0, 2, 1)
+        return tvm.topi.testing.batch_matmul(a, b)
+
+    def compiling(f, name="test_batch_matmul", ext=".so"):
+        path = name + ext
+        f.export_library(path)
+        mod = tvm.runtime.load_module(path)
+        f = mod[name]
+        return f
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func(lib.__name__ + ".matmul", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        name = "test_batch_matmul"
+        f = tvm.compile(
+            te.create_prim_func([input1_data, input2_data, final_result]), target=target
+        )
+        if target == "c":
+            f = compiling(f, name)
+        matrix_input1 = tvm.runtime.tensor(
+            np.random.uniform(size=ashape).astype(input1_data.dtype), dev
+        )
+        matrix_input2 = tvm.runtime.tensor(
+            np.random.uniform(size=bshape).astype(input2_data.dtype), dev
+        )
+        matrix_result = tvm.runtime.tensor(
+            np.zeros((batch, matrix_n, matrix_m), dtype=final_result.dtype), dev
+        )
+        f(matrix_input1, matrix_input2, matrix_result)
+        tvm.testing.assert_allclose(
+            matrix_result.numpy(),
+            get_numpy(matrix_input1.numpy(), matrix_input2.numpy(), transa, transb),
+            rtol=1e-5,
+        )
+
+    verify("llvm")
+    verify("c")
+
+
+def test_batch_matmul():
+    """Tests of matmul op where matrices are in batch"""
+    verify_batch_matmul(16, 16, 235, 128, 1024, cblas)
+    verify_batch_matmul(16, 16, 235, 128, 1024, cblas, True, False)
+    verify_batch_matmul(16, 16, 235, 128, 1024, cblas, False, True)
+    verify_batch_matmul(16, 16, 235, 128, 1024, cblas, True, True)
+    verify_batch_matmul(16, 16, 235, 128, 1024, mkl)
+    verify_batch_matmul(16, 16, 235, 128, 1024, mkl, True, False)
+    verify_batch_matmul(16, 16, 235, 128, 1024, mkl, False, True)
+    verify_batch_matmul(16, 16, 235, 128, 1024, mkl, True, True)
+    verify_batch_matmul(16, 1, 235, 128, 1024, cblas)
+    verify_batch_matmul(1, 16, 235, 128, 1024, cblas)
+    verify_batch_matmul(16, 1, 235, 128, 1024, cblas)
+    verify_batch_matmul(1, 16, 235, 128, 1024, cblas)
+    verify_batch_matmul(16, 1, 235, 128, 1024, mkl)
+    verify_batch_matmul(1, 16, 235, 128, 1024, mkl)
+    verify_batch_matmul(16, 1, 235, 128, 1024, mkl)
+    verify_batch_matmul(1, 16, 235, 128, 1024, mkl)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas, True, False)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas, False, False)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas, True, True)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl, True, False)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl, False, False)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl, True, True)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl)
+
+
+if __name__ == "__main__":
+    test_matmul_add()
+    test_quantized_matmul_add()
+    test_batch_matmul()

tests/python/contrib/test_coreml_runtime.py

@@ -0,0 +1,108 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+import os
+
+import numpy as np
+import pytest
+
+import tvm
+from tvm import rpc, te
+from tvm.contrib import coreml_runtime
+from tvm.support import utils, xcode
+
+proxy_host = os.environ.get("TVM_IOS_RPC_PROXY_HOST", "127.0.0.1")
+proxy_port = os.environ.get("TVM_IOS_RPC_PROXY_PORT", 9090)
+destination = os.environ.get("TVM_IOS_RPC_DESTINATION", "")
+key = "iphone"
+
+
+@pytest.mark.skip("skip because coremltools is not available in CI")
+def test_coreml_runtime():
+    import coremltools
+    from coremltools.models.neural_network import NeuralNetworkBuilder
+
+    def create_coreml_model():
+        shape = (2,)
+        alpha = 2
+
+        inputs = [
+            ("input0", coremltools.models.datatypes.Array(*shape)),
+            ("input1", coremltools.models.datatypes.Array(*shape)),
+        ]
+        outputs = [
+            ("output0", coremltools.models.datatypes.Array(*shape)),
+            ("output1", coremltools.models.datatypes.Array(*shape)),
+        ]
+        builder = NeuralNetworkBuilder(inputs, outputs)
+        builder.add_elementwise(
+            name="Add", input_names=["input0", "input1"], output_name="output0", mode="ADD"
+        )
+        builder.add_elementwise(
+            name="Mul", alpha=alpha, input_names=["input0"], output_name="output1", mode="MULTIPLY"
+        )
+        return coremltools.models.MLModel(builder.spec)
+
+    def verify(coreml_model, model_path, dev):
+        coreml_model = create_coreml_model()
+
+        out_spec = coreml_model.output_description._fd_spec
+        out_names = [spec.name for spec in out_spec]
+
+        # inference via coremltools
+        inputs = {}
+        for in_spec in coreml_model.input_description._fd_spec:
+            name = in_spec.name
+            shape = in_spec.type.multiArrayType.shape
+            inputs[name] = np.random.random_sample(shape)
+
+        coreml_outputs = [coreml_model.predict(inputs)[name] for name in out_names]
+
+        # inference via tvm coreml runtime
+        runtime = coreml_runtime.create("main", model_path, dev)
+        for name in inputs:
+            runtime.set_input(name, tvm.runtime.tensor(inputs[name], dev))
+        runtime.invoke()
+        tvm_outputs = [runtime.get_output(i).numpy() for i in range(runtime.get_num_outputs())]
+
+        for c_out, t_out in zip(coreml_outputs, tvm_outputs):
+            np.testing.assert_almost_equal(c_out, t_out, 3)
+
+    def check_remote(coreml_model):
+        temp = utils.tempdir()
+        compiled_model = xcode.compile_coreml(coreml_model, out_dir=temp.temp_dir)
+        xcode.popen_test_rpc(
+            proxy_host, proxy_port, key, destination=destination, libs=[compiled_model]
+        )
+        compiled_model = os.path.basename(compiled_model)
+        remote = rpc.connect(proxy_host, proxy_port, key=key)
+        dev = remote.cpu(0)
+        verify(coreml_model, compiled_model, dev)
+
+    def check_local(coreml_model):
+        temp = utils.tempdir()
+        compiled_model = xcode.compile_coreml(coreml_model, out_dir=temp.temp_dir)
+        dev = tvm.cpu(0)
+        verify(coreml_model, compiled_model, dev)
+
+    coreml_model = create_coreml_model()
+    check_remote(coreml_model)
+    check_local(coreml_model)
+
+
+if __name__ == "__main__":
+    test_coreml_runtime()

tests/python/contrib/test_cutlass_gemm.py

@@ -0,0 +1,399 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import ml_dtypes
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.contrib.pickle_memoize import memoize
+from tvm.testing import env
+
+
+def get_random_tensor(shape, dtype):
+    if dtype == "int8":
+        return np.random.randint(-128, 128, shape).astype(dtype)
+    elif dtype == "uint8":
+        return np.random.randint(0, 256, shape).astype(dtype)
+    return np.random.uniform(-1, 1, shape).astype(dtype)
+
+
+def verify_group_gemm(
+    func_name, M, N, K, num_groups, x_dtype, weight_dtype, out_dtype, use_scale, rtol, atol
+):
+    group_gemm_func = tvm.get_global_func(func_name, allow_missing=True)
+    if group_gemm_func is None:
+        print(f"Skipped as {func_name} is not available")
+        return
+
+    @memoize("tvm.contrib.cutlass.test_group_gemm_sm90")
+    def get_ref_data():
+        assert M % num_groups == 0
+        M_per_group = M // num_groups
+        a_np = get_random_tensor((M, K), x_dtype)
+        b_np = get_random_tensor((num_groups, N, K), weight_dtype)
+        indptr_np = np.arange(1, num_groups + 1).astype("int64") * M_per_group
+        c_np = np.concatenate(
+            [a_np[i * M_per_group : (i + 1) * M_per_group] @ b_np[i].T for i in range(num_groups)],
+            axis=0,
+        )
+        return a_np, b_np, indptr_np, c_np
+
+    def to_numpy_dtype(dtype):
+        mapping = {"float8_e5m2": ml_dtypes.float8_e5m2, "float8_e4m3fn": ml_dtypes.float8_e4m3fn}
+        return mapping.get(dtype, dtype)
+
+    a_np, b_np, indptr_np, c_np = get_ref_data()
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        a_nd = tvm.runtime.tensor(a_np.astype(to_numpy_dtype(x_dtype)), device=dev)
+        b_nd = tvm.runtime.tensor(b_np.astype(to_numpy_dtype(weight_dtype)), device=dev)
+        c_nd = tvm.runtime.empty(c_np.shape, dtype=out_dtype, device=dev)
+        indptr_nd = tvm.runtime.tensor(indptr_np, device=dev)
+        workspace = tvm.runtime.empty((4096 * 1024,), dtype="uint8", device=dev)
+        if use_scale:
+            scale = tvm.runtime.tensor(np.array([1.0], dtype="float32"), device=dev)
+            group_gemm_func(a_nd, b_nd, indptr_nd, workspace, scale, c_nd)
+        else:
+            group_gemm_func(a_nd, b_nd, indptr_nd, workspace, c_nd)
+        tvm.testing.assert_allclose(c_nd.numpy(), c_np, rtol=rtol, atol=atol)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_CUTLASS"), reason="need cutlass")
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_group_gemm_sm90():
+    verify_group_gemm(
+        "cutlass.group_gemm",
+        8,
+        128,
+        128,
+        4,
+        "float16",
+        "float16",
+        "float16",
+        False,
+        rtol=1e-3,
+        atol=1e-3,
+    )
+    verify_group_gemm(
+        "cutlass.group_gemm_e5m2_e5m2_fp16",
+        8,
+        16,
+        16,
+        4,
+        "float8_e5m2",
+        "float8_e5m2",
+        "float16",
+        True,
+        rtol=1e-1,
+        atol=1,
+    )
+    verify_group_gemm(
+        "cutlass.group_gemm_e4m3_e4m3_fp16",
+        8,
+        16,
+        16,
+        4,
+        "float8_e4m3fn",
+        "float8_e4m3fn",
+        "float16",
+        True,
+        rtol=1e-1,
+        atol=1,
+    )
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_CUTLASS"), reason="need cutlass")
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_group_gemm_sm100():
+    verify_group_gemm(
+        "cutlass.group_gemm",
+        8,
+        128,
+        128,
+        4,
+        "bfloat16",
+        "bfloat16",
+        "bfloat16",
+        False,
+        rtol=1e-2,
+        atol=1e-3,
+    )
+
+
+def rowwise_quant_fp8_e4m3(shape: tuple[int, int], block_size: tuple[int, int], dtype: str):
+    x_full_np = (np.random.rand(*shape) * 2 - 1).astype(dtype)
+    x_scale_shape = (
+        *shape[:-1],
+        (shape[-1] + block_size[1] - 1) // block_size[1],
+    )
+    # For each (block_size[1]) block, compute the max abs value of `w_full_np`
+    x_max_abs_np = np.zeros(x_scale_shape, dtype="float32")
+    for i in range(x_scale_shape[-1]):
+        x_max_abs_np[..., i] = np.max(
+            np.abs(x_full_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])]),
+            axis=-1,
+        )[0]
+    # Scale is the `x_max_abs_np` divided by the max value of quant_dtype in ml_dtypes
+    fp8_max = float(ml_dtypes.finfo("float8_e4m3fn").max)
+    x_scale_np = x_max_abs_np / fp8_max
+    # `x_np` is the `x_full_np` divided by the `x_scale_np` (with block awareness),
+    # clamped to (-fp8_max, fp8_max), and cast to `quant_dtype`
+    x_np = np.zeros_like(x_full_np, dtype="float8_e4m3fn")
+    for i in range(x_scale_shape[-1]):
+        x_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])] = np.clip(
+            x_full_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])]
+            / x_scale_np[..., i : i + 1],
+            -fp8_max,
+            fp8_max,
+        )
+
+    x_scale_np = np.random.rand(*x_scale_np.shape).astype("float32") / fp8_max
+    for i in range(x_scale_shape[-1]):
+        x_full_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])] = (
+            x_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])].astype(
+                x_scale_np.dtype
+            )
+            * x_scale_np[..., i : i + 1]
+        )
+    return x_np, x_scale_np
+
+
+def blockwise_quant_fp8_e4m3(shape: tuple[int, int], block_size: tuple[int, int], dtype: str):
+    w_full_np = (np.random.rand(*shape) * 2 - 1).astype(dtype)
+    w_scale_shape = (
+        *shape[:-2],
+        (shape[-2] + block_size[0] - 1) // block_size[0],
+        (shape[-1] + block_size[1] - 1) // block_size[1],
+    )
+    # For each (block_size[0], block_size[1]) block, compute the max abs value of `w_full_np`
+    w_max_abs_np = np.zeros(w_scale_shape, dtype="float32")
+    for i in range(w_scale_shape[-2]):
+        for j in range(w_scale_shape[-1]):
+            block_shape = (
+                *shape[:-2],
+                min(block_size[0], shape[-2] - i * block_size[0]),
+                min(block_size[1], shape[-1] - j * block_size[1]),
+            )
+            w_max_abs_np[..., i, j] = np.max(
+                np.abs(
+                    w_full_np[
+                        ...,
+                        i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                        j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                    ]
+                ).reshape(*shape[:-2], block_shape[-2] * block_shape[-1]),
+                axis=-1,
+            )
+    # Scale is the `w_max_abs_np` divided by the max value of quant_dtype in ml_dtypes
+    fp8_max = float(ml_dtypes.finfo("float8_e4m3fn").max)
+    w_scale_np = w_max_abs_np / fp8_max
+    # `w_np` is the `w_full_np` divided by the `w_scale_np` (with block awareness),
+    # clamped to (-fp8_max, fp8_max), and cast to `quant_dtype`
+    w_np = np.zeros_like(w_full_np, dtype="float8_e4m3fn")
+    if len(w_scale_shape) == 2:
+        for i in range(w_scale_shape[-2]):
+            for j in range(w_scale_shape[-1]):
+                w_np[
+                    i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                    j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                ] = np.clip(
+                    w_full_np[
+                        i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                        j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                    ]
+                    / w_scale_np[..., i, j],
+                    -fp8_max,
+                    fp8_max,
+                )
+    else:
+        for e in range(w_scale_shape[0]):
+            for i in range(w_scale_shape[-2]):
+                for j in range(w_scale_shape[-1]):
+                    w_np[
+                        e,
+                        i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                        j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                    ] = np.clip(
+                        w_full_np[
+                            e,
+                            i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                            j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                        ]
+                        / w_scale_np[e, i, j],
+                        -fp8_max,
+                        fp8_max,
+                    )
+
+    w_scale_np = np.random.rand(*w_scale_np.shape).astype("float32") / fp8_max
+    return w_np, w_scale_np
+
+
+def blockwise_matmul(
+    x_fp8_np: np.ndarray,
+    x_scale_np: np.ndarray,
+    w_np: np.ndarray,
+    w_scale_np: np.ndarray,
+    block_size: tuple[int, int],
+    dtype: str,
+):
+    o_np = np.zeros((x_fp8_np.shape[0], w_np.shape[0]), dtype=dtype)
+    for j in range(w_scale_np.shape[0]):
+        for k in range(w_scale_np.shape[1]):
+            o_np[:, j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[0])] += (
+                np.matmul(
+                    x_fp8_np[
+                        :, k * block_size[1] : min((k + 1) * block_size[1], x_fp8_np.shape[1])
+                    ].astype(dtype),
+                    w_np[
+                        j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[0]),
+                        k * block_size[1] : min((k + 1) * block_size[1], w_np.shape[1]),
+                    ].T.astype(dtype),
+                )
+                * x_scale_np[:, k : k + 1]
+                * w_scale_np[j, k]
+            )
+    return o_np
+
+
+def blockwise_bmm(
+    x_fp8_np: np.ndarray,
+    x_scale_np: np.ndarray,
+    w_np: np.ndarray,
+    w_scale_np: np.ndarray,
+    block_size: tuple[int, int],
+    dtype: str,
+):
+    o_np = np.zeros((x_fp8_np.shape[0], x_fp8_np.shape[1], w_np.shape[1]), dtype=dtype)
+    for j in range(w_scale_np.shape[1]):
+        for k in range(w_scale_np.shape[2]):
+            o_np[..., j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[1])] += (
+                np.matmul(
+                    x_fp8_np[
+                        ..., k * block_size[1] : min((k + 1) * block_size[1], x_fp8_np.shape[2])
+                    ].astype(dtype),
+                    w_np[
+                        ...,
+                        j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[1]),
+                        k * block_size[1] : min((k + 1) * block_size[1], w_np.shape[2]),
+                    ]
+                    .transpose(0, 2, 1)
+                    .astype(dtype),
+                )
+                * x_scale_np[..., k : k + 1]
+                * w_scale_np[..., j : j + 1, k : k + 1]
+            )
+    return o_np
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_CUTLASS"), reason="need cutlass")
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_fp8_e4m3_groupwise_scaled_gemm():
+    M = 16
+    N = 4608
+    K = 896
+    block_size = (128, 128)
+    assert N % 128 == 0 and K % 128 == 0  # Only support N/K are multiple of 128
+
+    func_name = "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn"
+    gemm_func = tvm.get_global_func(func_name, allow_missing=True)
+    if gemm_func is None:
+        print(f"Skipped as {func_name} is not available")
+        return
+
+    dtype = "bfloat16"
+    x_np, x_scale_np = rowwise_quant_fp8_e4m3((M, K), block_size, dtype)
+    w_np, w_scale_np = blockwise_quant_fp8_e4m3((N, K), block_size, dtype)
+    o_np = blockwise_matmul(x_np, x_scale_np, w_np, w_scale_np, block_size, dtype)
+
+    def run_and_check():
+        device = tvm.cuda(0)
+        x_tvm = tvm.runtime.tensor(x_np, device=device)
+        x_scale_tvm = tvm.runtime.tensor(x_scale_np.T, device=device)
+        w_tvm = tvm.runtime.tensor(w_np, device=device)
+        w_scale_tvm = tvm.runtime.tensor(w_scale_np, device=device)
+        workspace = tvm.runtime.empty((4096 * 1024,), dtype="uint8", device=device)
+        o_tvm = tvm.runtime.empty((M, N), dtype=dtype, device=device)
+        gemm_func(
+            x_tvm,
+            w_tvm,
+            x_scale_tvm,
+            w_scale_tvm,
+            workspace,
+            block_size[0],
+            block_size[1],
+            o_tvm,
+        )
+        tvm.testing.assert_allclose(o_tvm.numpy(), o_np, rtol=1e-4, atol=0.5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_CUTLASS"), reason="need cutlass")
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_fp8_e4m3_groupwise_scaled_bmm():
+    B = 16
+    M = 40
+    N = 512
+    K = 128
+    block_size = (128, 128)
+    assert N % 128 == 0 and K % 128 == 0  # Only support N/K are multiple of 128
+
+    func_name = "cutlass.groupwise_scaled_bmm_e4m3fn_e4m3fn"
+    gemm_func = tvm.get_global_func(func_name, allow_missing=True)
+    if gemm_func is None:
+        print(f"Skipped as {func_name} is not available")
+        return
+
+    dtype = "bfloat16"
+    x_np, x_scale_np = rowwise_quant_fp8_e4m3((B, M, K), block_size, dtype)
+    w_np, w_scale_np = blockwise_quant_fp8_e4m3((B, N, K), block_size, dtype)
+    o_np = blockwise_bmm(x_np, x_scale_np, w_np, w_scale_np, block_size, dtype)
+
+    def run_and_check():
+        device = tvm.cuda(0)
+        x_tvm = tvm.runtime.tensor(x_np, device=device)
+        x_scale_tvm = tvm.runtime.tensor(x_scale_np.transpose(0, 2, 1), device=device)
+        w_tvm = tvm.runtime.tensor(w_np, device=device)
+        w_scale_tvm = tvm.runtime.tensor(w_scale_np, device=device)
+        workspace = tvm.runtime.empty((4096 * 1024,), dtype="uint8", device=device)
+        o_tvm = tvm.runtime.empty((B, M, N), dtype=dtype, device=device)
+        gemm_func(
+            x_tvm,
+            w_tvm,
+            x_scale_tvm,
+            w_scale_tvm,
+            workspace,
+            block_size[0],
+            block_size[1],
+            o_tvm,
+        )
+        tvm.testing.assert_allclose(o_tvm.numpy(), o_np, rtol=1e-4, atol=0.5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/contrib/test_dlpack.py

@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm import te
+from tvm.contrib.dlpack import to_pytorch_func
+
+
+def verify_torch_dlpack():
+    a = np.random.randn(1337)
+    tvm_a = tvm.runtime.tensor(a)
+    np.testing.assert_equal(tvm.runtime.from_dlpack(tvm_a).numpy(), a)
+
+    try:
+        import torch
+        import torch.utils.dlpack
+
+        x = torch.rand(56, 56)
+        tvm_x = tvm.runtime.from_dlpack(torch.utils.dlpack.to_dlpack(x))
+        np.testing.assert_equal(x.numpy(), tvm_x.numpy())
+        y = tvm.runtime.from_dlpack(tvm_x)
+        np.testing.assert_equal(y.numpy(), tvm_x.numpy())
+        np.testing.assert_equal(torch.utils.dlpack.from_dlpack(y).numpy(), tvm_x.numpy())
+
+        n = tvm.runtime.convert(137)
+        xx = torch.rand(137, 137)
+        yy = torch.rand(137, 137)
+        zz2 = torch.empty(137, 137)
+        zz = xx.mm(yy)
+        XX = te.placeholder((n, n), name="X")
+        YY = te.placeholder((n, n), name="Y")
+
+        k = te.reduce_axis((0, n), name="k")
+        ZZ = te.compute((n, n), lambda i, j: te.sum(XX[i, k] * YY[k, j], axis=k))
+        # No need to speficy target_host if it's llvm
+        # Otherwise you will need to specify the target and target_host
+        f = tvm.compile(te.create_prim_func([XX, YY, ZZ]))
+
+        f_pytorch = to_pytorch_func(f)
+        zz2 = torch.empty(137, 137)
+        f_pytorch(xx, yy, zz2)
+        tvm.testing.assert_allclose(zz.numpy(), zz2.numpy(), rtol=1e-4, atol=1e-4)
+
+    except ImportError:
+        pass
+
+
+def test_torch_dlpack():
+    # Run dlpack interoperability test a few times to make sure it's stable.
+    for i in range(5):
+        verify_torch_dlpack()
+
+
+if __name__ == "__main__":
+    test_torch_dlpack()

tests/python/contrib/test_example_npu.py

@@ -0,0 +1,277 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Tests for Example NPU Backend
+
+This test file demonstrates how to test a custom NPU backend
+implementation using TVM's testing infrastructure.
+"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
+from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
+from tvm.script import relax as R
+
+
+@tvm.script.ir_module
+class MatmulReLU:
+    """Example module with matrix multiplication and ReLU"""
+
+    @R.function
+    def main(
+        x: R.Tensor((2, 4), "float32"),
+        w: R.Tensor((4, 8), "float32"),
+    ) -> R.Tensor((2, 8), "float32"):
+        with R.dataflow():
+            y = relax.op.matmul(x, w)
+            z = relax.op.nn.relu(y)
+            R.output(z)
+        return z
+
+
+@tvm.script.ir_module
+class Conv2dReLU:
+    """Example module with 2D convolution and ReLU"""
+
+    @R.function
+    def main(
+        x: R.Tensor((1, 3, 32, 32), "float32"),
+        w: R.Tensor((16, 3, 3, 3), "float32"),
+    ) -> R.Tensor((1, 16, 30, 30), "float32"):
+        with R.dataflow():
+            y = relax.op.nn.conv2d(x, w)
+            z = relax.op.nn.relu(y)
+            R.output(z)
+        return z
+
+
+@tvm.script.ir_module
+class MultipleOps:
+    """Example module with multiple operations that can be fused"""
+
+    @R.function
+    def main(
+        x: R.Tensor((1, 16, 32, 32), "float32"),
+    ) -> R.Tensor((1, 16, 16, 16), "float32"):
+        with R.dataflow():
+            # First ReLU
+            y = relax.op.nn.relu(x)
+            # Max pooling
+            z = relax.op.nn.max_pool2d(y, pool_size=(2, 2), strides=(2, 2))
+            # Second ReLU
+            out = relax.op.nn.relu(z)
+            R.output(out)
+        return out
+
+
+@tvm.script.ir_module
+class Softmax:
+    """Example module with softmax"""
+
+    @R.function
+    def main(x: R.Tensor((2, 8), "float32")) -> R.Tensor((2, 8), "float32"):
+        with R.dataflow():
+            z = relax.op.nn.softmax(x)
+            R.output(z)
+        return z
+
+
+# Check if the example NPU runtime is available
+has_example_npu_codegen = tvm.get_global_func("relax.ext.example_npu", True)
+has_example_npu_runtime = tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
+has_example_npu = has_example_npu_codegen and has_example_npu_runtime
+
+example_npu_enabled = pytest.mark.skipif(
+    not has_example_npu,
+    reason="Example NPU backend not enabled. Compile with the example NPU runtime.",
+)
+
+
+def test_example_npu_patterns_registered():
+    """Test that all expected patterns are registered"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    patterns = get_patterns_with_prefix("example_npu")
+    pattern_names = {p.name for p in patterns}
+
+    # Core patterns that should always be available
+    core_patterns = {
+        "example_npu.dense",
+        "example_npu.matmul",
+        "example_npu.conv1d",
+        "example_npu.conv2d",
+        "example_npu.max_pool2d",
+    }
+
+    assert core_patterns.issubset(pattern_names), (
+        f"Missing core patterns: {core_patterns - pattern_names}"
+    )
+
+    # Check that at least some activation patterns are available
+    activation_patterns = {name for name in pattern_names if "relu" in name or "sigmoid" in name}
+    assert len(activation_patterns) > 0, "No activation patterns found"
+
+
+@example_npu_enabled
+def test_example_npu_matmul_relu_partitioning():
+    """Test graph partitioning for MatMul + ReLU pattern"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = MatmulReLU
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Partition the graph
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+
+    # Verify partitioning happened
+    assert partitioned_mod is not None
+
+    # Check that composite functions were created
+    for gvar, func in partitioned_mod.functions.items():
+        if gvar.name_hint != "main":
+            # This should be a composite function
+            assert "Composite" in str(func)
+
+
+@example_npu_enabled
+def test_example_npu_conv2d_relu_partitioning():
+    """Test graph partitioning for Conv2D + ReLU pattern"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = Conv2dReLU
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Partition the graph
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+
+    assert partitioned_mod is not None
+
+
+@example_npu_enabled
+def test_example_npu_multiple_ops():
+    """Test partitioning with multiple fusable operations"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = MultipleOps
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Partition the graph
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+
+    assert partitioned_mod is not None
+
+
+@example_npu_enabled
+def test_example_npu_softmax_partitioning():
+    """Test graph partitioning for softmax pattern"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = Softmax
+    patterns = get_patterns_with_prefix("example_npu")
+
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+
+    assert partitioned_mod is not None
+
+    for gvar, func in partitioned_mod.functions.items():
+        if gvar.name_hint != "main":
+            assert "Composite" in str(func)
+
+
+@example_npu_enabled
+def test_example_npu_codegen():
+    """Test code generation for the example NPU backend"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = MatmulReLU
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Partition and generate code
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+    partitioned_mod = RunCodegen()(partitioned_mod)
+
+    assert partitioned_mod is not None
+
+    # The module should now contain external function calls
+    main_func = partitioned_mod["main"]
+    assert main_func is not None
+
+
+@example_npu_enabled
+def test_example_npu_runtime_execution():
+    """Test end-to-end execution with the example NPU runtime"""
+    import tvm.relax.backend.contrib.example_npu
+
+    # Create simple test inputs
+    np.random.seed(42)
+    x_np = np.random.randn(2, 4).astype("float32")
+    w_np = np.random.randn(4, 8).astype("float32")
+
+    # Expected output (computed with NumPy)
+    expected = np.maximum(0, np.matmul(x_np, w_np))
+
+    # Build and run with example NPU backend
+    mod = MatmulReLU
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Apply transformations
+    mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
+    mod = MergeCompositeFunctions()(mod)
+    mod = RunCodegen()(mod)
+
+    # Build the module
+    target = tvm.target.Target("llvm")
+    with tvm.transform.PassContext(opt_level=3):
+        built = relax.build(mod, target)
+
+    # Create VM and run
+    vm = relax.VirtualMachine(built, tvm.cpu())
+
+    x_tvm = tvm.runtime.tensor(x_np, tvm.cpu())
+    w_tvm = tvm.runtime.tensor(w_np, tvm.cpu())
+
+    result = vm["main"](x_tvm, w_tvm)
+
+    # Verify the result shape is correct (the runtime is a stub and does not compute numerically)
+    assert result.numpy().shape == expected.shape
+
+
+if __name__ == "__main__":
+    # Run tests locally for debugging
+    test_example_npu_patterns_registered()
+
+    if has_example_npu:
+        print("Example NPU backend is available, running tests...")
+        test_example_npu_matmul_relu_partitioning()
+        test_example_npu_conv2d_relu_partitioning()
+        test_example_npu_softmax_partitioning()
+        test_example_npu_multiple_ops()
+        test_example_npu_codegen()
+        test_example_npu_runtime_execution()
+        print("All tests passed!")
+    else:
+        print("Example NPU backend not available. Compile with example NPU runtime to run tests.")

tests/python/contrib/test_hipblas.py

@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E741
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te
+from tvm.contrib import hipblas
+from tvm.testing import env
+
+
+def verify_matmul_add(in_dtype, out_dtype, rtol=1e-5):
+    n = 1024
+    l = 128
+    m = 236
+    A = te.placeholder((n, l), name="A", dtype=in_dtype)
+    B = te.placeholder((l, m), name="B", dtype=in_dtype)
+    C = hipblas.matmul(A, B, dtype=out_dtype)
+
+    def verify(target="rocm"):
+        if not tvm.get_global_func("tvm.contrib.hipblas.matmul", True):
+            print("skip because extern function is not available")
+            return
+        f = tvm.compile(te.create_prim_func([A, B, C]), target=target)
+
+        def run_and_check():
+            dev = tvm.rocm(0)
+            a = tvm.runtime.tensor(np.random.uniform(0, 128, size=(n, l)).astype(A.dtype), dev)
+            b = tvm.runtime.tensor(np.random.uniform(0, 128, size=(l, m)).astype(B.dtype), dev)
+            c = tvm.runtime.tensor(np.zeros((n, m), dtype=C.dtype), dev)
+            f(a, b, c)
+            tvm.testing.assert_allclose(
+                c.numpy(),
+                np.dot(a.numpy().astype(C.dtype), b.numpy().astype(C.dtype)),
+                rtol=rtol,
+            )
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    verify()
+
+
+def roundoff(v, d):
+    return int(np.floor((v + d - 1) / d) * d)
+
+
+def verify_batch_matmul(Ashape, Bshape, Cshape, in_dtype, out_dtype, rtol=1e-5):
+    A = te.placeholder(Ashape, name="A", dtype=in_dtype)
+    B = te.placeholder(Bshape, name="B", dtype=in_dtype)
+    C = hipblas.batch_matmul(A, B, dtype=out_dtype)
+
+    f = tvm.compile(te.create_prim_func([A, B, C]), target="rocm")
+
+    def run_and_check():
+        dev = tvm.rocm(0)
+        if "int" in in_dtype:
+            a = tvm.runtime.tensor(np.random.uniform(1, 10, size=Ashape).astype(in_dtype), dev)
+            b = tvm.runtime.tensor(np.random.uniform(1, 10, size=Bshape).astype(in_dtype), dev)
+        else:
+            a = tvm.runtime.tensor(np.random.uniform(size=Ashape).astype(A.dtype), dev)
+            b = tvm.runtime.tensor(np.random.uniform(size=Bshape).astype(B.dtype), dev)
+
+        c = tvm.runtime.tensor(np.zeros(Cshape, dtype=C.dtype), dev)
+        f(a, b, c)
+        tvm.testing.assert_allclose(
+            c.numpy(),
+            np.matmul(a.numpy().astype(C.dtype), b.numpy().astype(C.dtype)).astype(C.dtype),
+            rtol=rtol,
+        )
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_matmul_add():
+    verify_matmul_add("float", "float", rtol=1e-3)
+    verify_matmul_add("float16", "float")
+    verify_matmul_add("float16", "float16", rtol=1e-2)
+    verify_matmul_add("int8", "int32")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_batch_matmul():
+    if not tvm.get_global_func("tvm.contrib.hipblas.batch_matmul", True):
+        print("skip because extern function is not available")
+        return
+
+    verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "float", "float")
+    verify_batch_matmul((16, 1024, 128), (1, 128, 236), (16, 1024, 236), "float", "float")
+    verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "float16", "float")
+    verify_batch_matmul((16, 1024, 128), (1, 128, 236), (16, 1024, 236), "float16", "float")
+    verify_batch_matmul(
+        (16, 1024, 128), (16, 128, 236), (16, 1024, 236), "float16", "float16", rtol=1e-2
+    )
+    verify_batch_matmul(
+        (16, 1024, 128), (1, 128, 236), (16, 1024, 236), "float16", "float16", rtol=1e-2
+    )
+
+    verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "int8", "int32")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/contrib/test_memoize.py

@@ -0,0 +1,128 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for tvm.contrib.pickle_memoize"""
+
+import os
+import pathlib
+import subprocess
+import sys
+import tempfile
+
+import tvm.testing
+
+TEST_SCRIPT_FILE = pathlib.Path(__file__).with_name("pickle_memoize_script.py").resolve()
+
+
+def test_cache_dir_not_in_current_working_dir():
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        temp_dir = pathlib.Path(temp_dir)
+        subprocess.check_call([sys.executable, str(TEST_SCRIPT_FILE), "1", "1"], cwd=temp_dir)
+
+        new_files = list(temp_dir.iterdir())
+        assert not new_files, (
+            "Use of tvm.contrib.pickle_memorize may not write to current directory."
+        )
+
+
+def test_current_directory_is_not_required_to_be_writable():
+    """TVM may be imported without directory permissions
+
+    This is a regression test.  In previous implementations, the
+    `tvm.contrib.pickle_memoize.memoize` function would write to the
+    current directory when importing TVM.  Import of a Python module
+    should not write to any directory.
+
+    """
+
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        temp_dir = pathlib.Path(temp_dir)
+
+        # User may read/cd into the temp dir, nobody may write to temp
+        # dir.
+        temp_dir.chmod(0o500)
+        subprocess.check_call([sys.executable, "-c", "import tvm"], cwd=temp_dir)
+
+
+def test_cache_dir_defaults_to_home_config_cache():
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        temp_dir = pathlib.Path(temp_dir)
+
+        subprocess.check_call([sys.executable, str(TEST_SCRIPT_FILE), "1", "0"], cwd=temp_dir)
+
+        new_files = list(temp_dir.iterdir())
+        assert not new_files, (
+            "Use of tvm.contrib.pickle_memorize may not write to current directory."
+        )
+
+        cache_dir = pathlib.Path.home().joinpath(".cache", "tvm", "pkl_memoize_py3")
+        assert cache_dir.exists()
+        cache_files = list(cache_dir.iterdir())
+        assert len(cache_files) >= 1
+
+
+def test_cache_dir_respects_xdg_cache_home():
+    with (
+        tempfile.TemporaryDirectory(prefix="tvm_") as temp_working_dir,
+        tempfile.TemporaryDirectory(prefix="tvm_") as temp_cache_dir,
+    ):
+        temp_cache_dir = pathlib.Path(temp_cache_dir)
+        temp_working_dir = pathlib.Path(temp_working_dir)
+
+        subprocess.check_call(
+            [sys.executable, str(TEST_SCRIPT_FILE), "1", "0"],
+            cwd=temp_working_dir,
+            env={
+                **os.environ,
+                "XDG_CACHE_HOME": temp_cache_dir.as_posix(),
+            },
+        )
+
+        new_files = list(temp_working_dir.iterdir())
+        assert not new_files, (
+            "Use of tvm.contrib.pickle_memorize may not write to current directory."
+        )
+
+        cache_dir = temp_cache_dir.joinpath("tvm", "pkl_memoize_py3")
+        assert cache_dir.exists()
+        cache_files = list(cache_dir.iterdir())
+        assert len(cache_files) == 1
+
+
+def test_cache_dir_only_created_when_used():
+    with (
+        tempfile.TemporaryDirectory(prefix="tvm_") as temp_working_dir,
+        tempfile.TemporaryDirectory(prefix="tvm_") as temp_cache_dir,
+    ):
+        temp_cache_dir = pathlib.Path(temp_cache_dir)
+        temp_working_dir = pathlib.Path(temp_working_dir)
+
+        subprocess.check_call(
+            [sys.executable, str(TEST_SCRIPT_FILE), "0", "1"],
+            cwd=temp_working_dir,
+            env={
+                **os.environ,
+                "XDG_CACHE_HOME": temp_cache_dir.as_posix(),
+            },
+        )
+
+        cache_dir = temp_cache_dir.joinpath("tvm", "pkl_memoize_py3")
+        assert not cache_dir.exists()
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/contrib/test_random.py

@@ -0,0 +1,200 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E722
+"""Configure pytest"""
+
+# pylint: disable=invalid-name
+import threading
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import rpc, te
+from tvm.contrib import random
+
+
+def test_randint():
+    """Tests randint function"""
+    m = 10240
+    n = 10240
+    A = random.randint(-127, 128, size=(m, n), dtype="int32")
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func("tvm.contrib.random.randint", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        f = tvm.compile(te.create_prim_func([A]), target=target)
+        a = tvm.runtime.tensor(np.zeros((m, n), dtype=A.dtype), dev)
+        f(a)
+        na = a.numpy()
+        assert abs(np.mean(na)) < 0.3
+        assert np.min(na) == -127
+        assert np.max(na) == 127
+
+    verify()
+
+
+def test_uniform():
+    """Tests uniform function"""
+    m = 10240
+    n = 10240
+    A = random.uniform(0, 1, size=(m, n))
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func("tvm.contrib.random.uniform", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        f = tvm.compile(te.create_prim_func([A]), target=target)
+        a = tvm.runtime.tensor(np.zeros((m, n), dtype=A.dtype), dev)
+        f(a)
+        na = a.numpy()
+        assert abs(np.mean(na) - 0.5) < 1e-1
+        assert abs(np.min(na) - 0.0) < 1e-3
+        assert abs(np.max(na) - 1.0) < 1e-3
+
+    verify()
+
+
+def test_normal():
+    """Tests normal function"""
+    m = 10240
+    n = 10240
+    A = random.normal(3, 4, size=(m, n))
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func("tvm.contrib.random.normal", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        f = tvm.compile(te.create_prim_func([A]), target=target)
+        a = tvm.runtime.tensor(np.zeros((m, n), dtype=A.dtype), dev)
+        f(a)
+        na = a.numpy()
+        assert abs(np.mean(na) - 3) < 1e-1
+        assert abs(np.std(na) - 4) < 1e-2
+
+    verify()
+
+
+@pytest.mark.gpu
+def test_random_fill():
+    """Tests random_fill function"""
+
+    def test_local(dev, dtype):
+        if not tvm.get_global_func("tvm.contrib.random.random_fill", True):
+            print("skip because extern function is not available")
+            return
+        value = tvm.runtime.empty((512, 512), dtype, dev)
+        random_fill = tvm.get_global_func("tvm.contrib.random.random_fill")
+        random_fill(value)
+
+        assert np.count_nonzero(value.numpy()) == 512 * 512
+
+        # make sure arithmentic doesn't overflow too
+        np_values = value.numpy()
+        assert np.isfinite(np_values * np_values + np_values).any()
+
+    def test_rpc(dtype):
+        if not tvm.get_global_func("tvm.contrib.random.random_fill", True):
+            print("skip because extern function is not available")
+            return
+        if not tvm.testing.device_enabled("rpc") or not tvm.runtime.enabled("llvm"):
+            return
+
+        def check_remote(server):
+            remote = rpc.connect(server.host, server.port)
+            value = tvm.runtime.empty((512, 512), dtype, remote.cpu())
+            random_fill = remote.get_function("tvm.contrib.random.random_fill")
+            random_fill(value)
+
+            assert np.count_nonzero(value.numpy()) == 512 * 512
+
+            # make sure arithmentic doesn't overflow too
+            np_values = value.numpy()
+            assert np.isfinite(np_values * np_values + np_values).any()
+
+        check_remote(rpc.Server("127.0.0.1"))
+
+    # Packed sub-byte dtypes (e.g. int4) are intentionally unsupported by
+    # random_fill since #19714 and raise an error instead.
+    for dtype in [
+        "bool",
+        "int8",
+        "uint8",
+        "int16",
+        "uint16",
+        "int32",
+        "int32",
+        "int64",
+        "uint64",
+        "float16",
+        "float32",
+        "float64",
+    ]:
+        for target, dev in tvm.testing.enabled_targets():
+            if tvm.target.Target(target).kind.name == "llvm":
+                test_local(dev, dtype)
+            else:
+                tvm.testing.run_with_gpu_lock(test_local, dev, dtype)
+        test_rpc(dtype)
+
+
+def test_random_fill_mt():
+    """Check random filler applicability in case of nontrivial thread pool configuration.
+    Particularly when MaxConcurrency != num_workers_used_ which is actual for big-little systems.
+    """
+    no_exception_happened = True
+
+    def test_body():
+        try:
+            num_thread_used = 1
+            configure_threads = tvm.get_global_func("runtime.config_threadpool")
+            configure_threads(1, num_thread_used)
+
+            test_input = tvm.runtime.empty((10, 10))
+            random_fill = tvm.get_global_func("tvm.contrib.random.random_fill_for_measure")
+            random_fill(test_input)
+        except:  # pylint: disable=bare-except
+            nonlocal no_exception_happened
+            no_exception_happened = False
+
+    # ThreadPool object is thread local. To eliminate effect on other test cases put it into thread
+    x = threading.Thread(target=test_body)
+    x.start()
+    x.join()
+    assert no_exception_happened
+
+
+if __name__ == "__main__":
+    test_randint()
+    test_uniform()
+    test_normal()
+    test_random_fill()
+    test_random_fill_mt()

tests/python/contrib/test_rpc_proxy.py

@@ -0,0 +1,69 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Configure pytest"""
+
+# pylint: disable=invalid-name
+import logging
+import multiprocessing
+import time
+
+import tvm
+from tvm import rpc
+
+
+def rpc_proxy_check():
+    """This is a simple test function for RPC Proxy
+
+    It is not included as pytests, because:
+    - It depends on tornado
+    - It relies on the fact that Proxy starts before client and server connects,
+      which is often the case but not always
+
+    User can directly run this script to verify correctness.
+    """
+
+    try:
+        # pylint: disable=import-outside-toplevel
+        from tvm.rpc import proxy
+
+        web_port = 8888
+        prox = proxy.Proxy("127.0.0.1", web_port=web_port)
+
+        def check():
+            if not tvm.runtime.enabled("rpc"):
+                return
+
+            server = multiprocessing.Process(
+                target=proxy.websocket_proxy_server, args=(f"ws://localhost:{web_port}/ws", "x1")
+            )
+            # Need to make sure that the connection start after proxy comes up
+            time.sleep(0.1)
+            server.deamon = True
+            server.start()
+            client = rpc.connect(prox.host, prox.port, key="x1")
+            f1 = client.get_function("testing.echo")
+            assert f1(10) == 10
+            assert f1("xyz") == "xyz"
+
+        check()
+    except ImportError:
+        print("Skipping because tornado is not avaliable...")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    rpc_proxy_check()

tests/python/contrib/test_rpc_tracker.py

@@ -0,0 +1,152 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Configure pytest"""
+
+# pylint: disable=invalid-name
+import logging
+import time
+
+import tvm
+from tvm import rpc
+
+
+def check_server_drop():
+    """test when server drops"""
+    try:
+        # pylint: disable=import-outside-toplevel
+        from tvm.rpc import base, proxy, tracker
+
+        # pylint: disable=import-outside-toplevel
+        from tvm.rpc.base import TrackerCode
+
+        @tvm.register_global_func("rpc.test2.addone")
+        def addone(x):
+            return x + 1
+
+        def _put(tclient, value):
+            base.sendjson(tclient._sock, value)
+            base.recvjson(tclient._sock)
+
+        tserver = tracker.Tracker("127.0.0.1", 8888)
+        tproxy = proxy.Proxy("127.0.0.1", 8881, tracker_addr=("127.0.0.1", tserver.port))
+        tclient = rpc.connect_tracker("127.0.0.1", tserver.port)
+
+        server0 = rpc.Server(
+            "127.0.0.1", port=9099, tracker_addr=("127.0.0.1", tserver.port), key="abc"
+        )
+        server1 = rpc.Server(
+            "127.0.0.1", port=9099, tracker_addr=("127.0.0.1", tserver.port), key="xyz"
+        )
+        server2 = rpc.Server("127.0.0.1", tproxy.port, is_proxy=True, key="xyz")
+        server3 = rpc.Server("127.0.0.1", tproxy.port, is_proxy=True, key="xyz1")
+
+        # Fault tolerence to un-handled requested value
+        _put(tclient, [TrackerCode.REQUEST, "abc", "", 1])
+        _put(tclient, [TrackerCode.REQUEST, "xyz1", "", 1])
+
+        # Fault tolerence to stale worker value
+        _put(tclient, [TrackerCode.PUT, "xyz", (server1.port, "abc")])
+        _put(tclient, [TrackerCode.PUT, "xyz", (server1.port, "abcxxx")])
+        _put(tclient, [TrackerCode.PUT, "xyz", (tproxy.port, "abcxxx11")])
+
+        # Fault tolerence server timeout
+        def check_timeout(timeout, sleeptime):
+            def myfunc(remote):
+                time.sleep(sleeptime)
+                f1 = remote.get_function("rpc.test2.addone")
+                assert f1(10) == 11
+
+            try:
+                tclient.request_and_run("xyz", myfunc, session_timeout=timeout)
+            except RuntimeError:
+                pass
+            print(tclient.text_summary())
+            try:
+                remote = tclient.request("xyz", priority=0, session_timeout=timeout)
+                remote2 = tclient.request("xyz", session_timeout=timeout)
+                time.sleep(sleeptime)
+                f1 = remote.get_function("rpc.test2.addone")
+                assert f1(10) == 11
+                f1 = remote2.get_function("rpc.test2.addone")
+                assert f1(10) == 11
+
+            except RuntimeError:
+                pass
+            remote3 = tclient.request("abc")
+            f1 = remote3.get_function("rpc.test2.addone")
+            assert f1(10) == 11
+            remote3 = tclient.request("xyz1")
+            f1 = remote3.get_function("rpc.test2.addone")
+            assert f1(10) == 11
+
+        check_timeout(0.01, 0.1)
+        check_timeout(2, 0)
+        tserver.terminate()
+        server0.terminate()
+        server1.terminate()
+        server2.terminate()
+        server3.terminate()
+        tproxy.terminate()
+    except ImportError:
+        print("Skip because tornado is not available")
+
+
+def check_tracker_rejects_oversized_msg_size():
+    """Tracker must reject an oversized msg_size header and close the connection
+    instead of buffering an unbounded amount of data on a single TCP connection.
+
+    Regression test for the unbounded buffer growth defect in
+    TCPEventHandler.on_message. See MAX_TRACKER_MSG_BYTES in tracker.py.
+    """
+    try:
+        # pylint: disable=import-outside-toplevel
+        import socket
+        import struct
+
+        from tvm.rpc import base, tracker
+
+        tserver = tracker.Tracker(port=9180, port_end=9290, silent=True)
+        try:
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.settimeout(5)
+            sock.connect(("127.0.0.1", tserver.port))
+            # complete the 4-byte magic handshake
+            sock.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
+            magic_reply = sock.recv(4)
+            assert struct.unpack("<i", magic_reply)[0] == base.RPC_TRACKER_MAGIC
+
+            # send an oversized msg_size header (2 GiB)
+            sock.sendall(struct.pack("<i", 0x7FFFFFFF))
+
+            # server must close the connection (no payload buffering)
+            for _ in range(20):
+                chunk = sock.recv(4096)
+                if chunk == b"":
+                    break
+                time.sleep(0.05)
+            else:
+                raise AssertionError("tracker did not close connection after oversized msg_size")
+        finally:
+            tserver.terminate()
+    except ImportError:
+        print("Skip because tornado is not available")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    check_server_drop()
+    check_tracker_rejects_oversized_msg_size()

tests/python/contrib/test_sort.py

@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E741
+"""Configure pytest"""
+
+# pylint: disable=invalid-name
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm import te
+
+
+def test_sort():
+    """Tests sort function"""
+    n = 2
+    l = 5
+    m = 3
+    data = te.placeholder((n, l, m), name="data")
+    sort_num = te.placeholder((n, m), name="sort_num", dtype="int32")
+    axis = 1
+    is_ascend = False
+    out = te.extern(
+        data.shape,
+        [data, sort_num],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.sort.argsort_nms", ins[0], ins[1], outs[0], axis, is_ascend
+        ),
+        dtype="int32",
+        name="sort_tensor",
+    )
+    input_data = [
+        [[1, 2, 3], [2, 4.5, 3.5], [1.1, 0.5, 1], [3.2, -5, 0.5], [1.5, 0, 0]],
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]],
+    ]
+    sort_num_input = [[1, 2, 3], [4, 5, 5]]
+    sorted_index = [
+        [[0, 1, 1], [1, 0, 0], [2, 2, 2], [3, 3, 3], [4, 4, 4]],
+        [[3, 4, 4], [2, 3, 3], [1, 2, 2], [0, 1, 1], [4, 0, 0]],
+    ]
+
+    dev = tvm.cpu(0)
+    target = "llvm"
+    f = tvm.compile(te.create_prim_func([data, sort_num, out]), target=target)
+    a = tvm.runtime.tensor(np.array(input_data).astype(data.dtype.dtype), dev)
+    b = tvm.runtime.tensor(np.array(sort_num_input).astype(sort_num.dtype.dtype), dev)
+    c = tvm.runtime.tensor(np.zeros(a.shape, dtype=out.dtype.dtype), dev)
+    f(a, b, c)
+    tvm.testing.assert_allclose(
+        c.numpy(), np.array(sorted_index).astype(out.dtype.dtype), rtol=1e-5
+    )
+
+
+def test_sort_np():
+    """Tests sort function using numpy"""
+    dshape = (1, 2, 3, 4, 5, 6)
+    axis = 4
+    reduced_shape = (1, 2, 3, 4, 6)
+    is_ascend = True
+    data = te.placeholder(dshape, name="data")
+    sort_num = te.placeholder(reduced_shape, name="sort_num", dtype="int32")
+    out = te.extern(
+        data.shape,
+        [data, sort_num],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.sort.argsort_nms", ins[0], ins[1], outs[0], axis, is_ascend
+        ),
+        dtype="int32",
+        name="sort_tensor",
+    )
+
+    dev = tvm.cpu(0)
+    target = "llvm"
+    f = tvm.compile(te.create_prim_func([data, sort_num, out]), target=target)
+
+    np_data = np.random.uniform(size=dshape)
+    np_out = np.argsort(np_data, axis=axis)
+    sort_num_input = np.full(reduced_shape, dshape[axis])
+    a = tvm.runtime.tensor(np.array(np_data).astype(data.dtype.dtype), dev)
+    b = tvm.runtime.tensor(np.array(sort_num_input).astype(sort_num.dtype.dtype), dev)
+    c = tvm.runtime.tensor(np.zeros(a.shape, dtype=out.dtype.dtype), dev)
+    f(a, b, c)
+    tvm.testing.assert_allclose(c.numpy(), np_out, rtol=1e-5)
+
+
+if __name__ == "__main__":
+    test_sort()
+    test_sort_np()

tests/python/contrib/test_tir_triton_integration.py

@@ -0,0 +1,134 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401
+import sys
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.relax.frontend import nn
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+from tvm.testing import env
+
+try:
+    import triton
+    import triton.language as tl
+    from packaging import version
+except ImportError:
+    pytestmark = pytest.skip("Triton is not available", allow_module_level=True)
+else:
+    if version.parse(triton.__version__) < version.parse("3.3.0"):
+        pytestmark = pytest.skip("Triton >= 3.3.0 is required", allow_module_level=True)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_tir_triton_integration():
+    @triton.jit
+    def add_kernel(
+        x_ptr,  # *Pointer* to first input vector.
+        y_ptr,  # *Pointer* to second input vector.
+        output_ptr,  # *Pointer* to output vector.
+        n_elements,  # Size of the vector.
+        BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.
+    ):
+        """Triton vector add kernel from its tutorial."""
+        pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
+        block_start = pid * BLOCK_SIZE
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+        x = tl.load(x_ptr + offsets, mask=mask)
+        y = tl.load(y_ptr + offsets, mask=mask)
+        output = x + y
+        tl.store(output_ptr + offsets, output, mask=mask)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def add(x_handle: T.handle, y_handle: T.handle, output_handle: T.handle) -> None:
+            T.func_attr({"global_symbol": "add"})
+            m = T.int64()
+            x = T.match_buffer(x_handle, (m,), "float32")
+            y = T.match_buffer(y_handle, (m,), "float32")
+            output = T.match_buffer(output_handle, (m,), "float32")
+            with T.sblock("root"):
+                T.reads(x[0:m], y[0:m])
+                T.writes(output[0:m])
+                BLOCK_SIZE = T.meta_var(64)
+                T.call_kernel(
+                    add_kernel,
+                    (T.ceildiv(m, BLOCK_SIZE),),
+                    x.data,
+                    y.data,
+                    output.data,
+                    m,
+                    BLOCK_SIZE,
+                    num_warps=8,
+                )
+
+        @R.function
+        def main(x: R.Tensor(("m",), "float32"), y: R.Tensor(("m",), "float32")):
+            m = T.int64()
+            with R.dataflow():
+                output = R.call_tir(Module.add, [x, y], relax.TensorType((m,), "float32"))
+                R.output(output)
+            return output
+
+    # Constexpr parameters (BLOCK_SIZE) stay in the kernel arguments, and the
+    # thread extent is 256 because the kernel is compiled with num_warps=8.
+    @I.ir_module(s_tir=True)
+    class Parsed:
+        @T.prim_func(s_tir=True)
+        def add(x_handle: T.handle, y_handle: T.handle, output_handle: T.handle):
+            m = T.int64()
+            x = T.match_buffer(x_handle, (m,))
+            y = T.match_buffer(y_handle, (m,))
+            output = T.match_buffer(output_handle, (m,))
+            with T.sblock("root"):
+                T.reads(x[0:m], y[0:m])
+                T.writes(output[0:m])
+                T.call_packed(
+                    "add_kernel",
+                    x.data,
+                    y.data,
+                    output.data,
+                    m,
+                    64,
+                    256,
+                    (m + T.int64(64) - T.int64(1)) // T.int64(64),
+                )
+
+    tvm.ir.assert_structural_equal(Module["add"], Parsed["add"])
+    assert len(Module.get_attr("external_mods")) == 1
+
+    with tvm.target.Target("cuda"):
+        lib = tvm.compile(Module)
+
+    def run_and_check():
+        device = tvm.cuda(0)
+        x_nd = tvm.runtime.tensor(np.random.rand(256).astype(np.float32), device)
+        y_nd = tvm.runtime.tensor(np.random.rand(256).astype(np.float32), device)
+        output_np = x_nd.numpy() + y_nd.numpy()
+        output_nd = tvm.runtime.vm.VirtualMachine(lib, device)["main"](x_nd, y_nd)
+        tvm.testing.assert_allclose(output_nd.numpy(), output_np, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)

tests/python/contrib/test_tvmjs.py

@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Test contrib.tvmjs"""
+
+import tempfile
+
+import numpy as np
+import pytest
+
+import tvm.testing
+from tvm.contrib import tvmjs
+
+dtype = tvm.testing.parameter(
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "float16",
+    "float32",
+    "float64",
+    "float8_e4m3fn",
+    "float8_e5m2",
+)
+
+
+def test_save_load_float8(dtype):
+    if "float8" in dtype or "bfloat16" in dtype:
+        ml_dtypes = pytest.importorskip("ml_dtypes")
+        np_dtype = np.dtype(getattr(ml_dtypes, dtype))
+    else:
+        np_dtype = np.dtype(dtype)
+
+    arr = np.arange(16, dtype=np_dtype)
+
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        tvmjs.dump_tensor_cache({"arr": arr}, temp_dir)
+        cache, _ = tvmjs.load_tensor_cache(temp_dir, tvm.cpu())
+
+    after_roundtrip = cache["arr"].numpy()
+
+    np.testing.assert_array_equal(arr, after_roundtrip)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/disco/test_callback.py

@@ -0,0 +1,175 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test sharded loader"""
+# pylint: disable=missing-docstring
+
+import pathlib
+import tempfile
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(tvm.runtime.disco is None, reason="disco runtime is not available")
+@pytest.mark.skipif(not env.has_nccl(), reason="need nccl")
+@pytest.mark.skipif(not env.has_multi_gpu(), reason="need multiple gpus")
+def test_callback():
+    """Simulate lazy loading of parameters in a callback
+
+    The output of a lazy parameter loading, which would accept a
+    callback to load the parameters.
+    """
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def slice_A(
+            A: T.Buffer((4, 4), "int32"),
+            rank: T.int64,
+            A_sharded: T.Buffer((2, 4), "int32"),
+        ):
+            for i, j in T.grid(2, 4):
+                with T.sblock("slice_A"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    A_sharded[vi, vj] = A[rank * 2 + vi, vj]
+
+        @T.prim_func(private=True, s_tir=True)
+        def slice_B(
+            B: T.Buffer((2, 2), "float32"),
+            rank: T.int64,
+            B_sharded: T.Buffer((2, 1), "float32"),
+        ):
+            for i in range(2):
+                with T.sblock("slice_B"):
+                    vi = T.axis.spatial(2, i)
+                    B_sharded[vi, 0] = B[vi, rank]
+
+        @R.function
+        def transform_params(
+            rank_arg: R.Prim("int64"),
+            fget_item: R.Callable([R.Any, R.Prim("int64")], R.Any),
+        ):
+            cls = Module
+
+            A = fget_item(R.str("A"), R.prim_value(0))
+            A = R.match_cast(A, R.Tensor([4, 4], "int32"))
+            A = R.call_tir(
+                cls.slice_A,
+                (A, rank_arg),
+                out_ty=R.Tensor([2, 4], "int32"),
+            )
+
+            B = fget_item(R.str("B"), R.prim_value(1))
+            B = R.match_cast(B, R.Tensor([2, 2], "float32"))
+            B = R.call_tir(
+                cls.slice_B,
+                (B, rank_arg),
+                out_ty=R.Tensor([2, 1], "float32"),
+            )
+
+            return (A, B)
+
+    pipeline = tvm.ir.transform.Sequential(
+        [
+            tvm.relax.transform.LegalizeOps(),
+            tvm.s_tir.dlight.ApplyDefaultSchedule(tvm.s_tir.dlight.gpu.Fallback()),
+        ],
+        name="pipeline",
+    )
+
+    with tvm.target.Target("cuda"):
+        mod = Module
+        mod = pipeline(mod)
+        built = tvm.compile(mod, "cuda")
+
+    num_shards = 2
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dir = pathlib.Path(temp_dir)
+
+        # TODO(Lunderberg): Update `disco.Session.load_vm_module` to
+        # allow a `tvm.runtime.Module` argument.  This would avoid the
+        # need for a temporary file.
+        shlib_path = temp_dir.joinpath("libtemp.so")
+        built.export_library(shlib_path)
+
+        def run_and_check():
+            session = tvm.runtime.disco.ProcessSession(num_workers=num_shards)
+            try:
+                session.import_python_module("tvm.exec.disco_worker")
+                session.init_ccl("nccl", *range(num_shards))
+
+                worker_device = session.get_global_func("runtime.disco.device")()
+                worker_id = session.get_global_func("runtime.disco.worker_rank")()
+                callback_maker = session.get_global_func("tests.disco.test_callback")
+                fget_item = callback_maker(worker_device)
+                vm = session.load_vm_module(shlib_path.as_posix())
+                transform_params = vm["transform_params"]
+
+                params = transform_params(worker_id, fget_item)
+
+                # Worker 0 is the same PID as the controlling scope, so
+                # `debug_get_from_remote(0)` returns the Tensor containing
+                # the output.
+                params_gpu0 = params.debug_get_from_remote(0)
+                assert params_gpu0[0].device == tvm.cuda(0)
+                assert params_gpu0[1].device == tvm.cuda(0)
+                np.testing.assert_array_equal(
+                    params_gpu0[0].numpy(),
+                    [
+                        [0, 1, 2, 3],
+                        [4, 5, 6, 7],
+                    ],
+                )
+                np.testing.assert_array_equal(
+                    params_gpu0[1].numpy(),
+                    [[0], [2]],
+                )
+
+                # Worker 1 is a different PID altogether, so
+                # `debug_get_from_remote(1)` returns a new Tensor within the
+                # calling scope's PID.
+                params_gpu1 = params.debug_get_from_remote(1)
+                assert params_gpu1[0].device == tvm.cpu()
+                assert params_gpu1[1].device == tvm.cpu()
+                np.testing.assert_array_equal(
+                    params_gpu1[0].numpy(),
+                    [
+                        [8, 9, 10, 11],
+                        [12, 13, 14, 15],
+                    ],
+                )
+                np.testing.assert_array_equal(
+                    params_gpu1[1].numpy(),
+                    [[1], [3]],
+                )
+            finally:
+                session.shutdown()
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()