From bafca9b22156d0ff00513f87ab5129190ba9a406 Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Sun, 13 Sep 2026 21:55:25 -0400 Subject: [PATCH 1/3] [TIR][ARITH] Add constraint visitor snapshots --- src/arith/constr_visitor.h | 554 +++++++++++++++++++++++++ tests/cpp/arith_constr_visitor_test.cc | 291 +++++++++++++ 2 files changed, 845 insertions(+) create mode 100644 src/arith/constr_visitor.h create mode 100644 tests/cpp/arith_constr_visitor_test.cc diff --git a/src/arith/constr_visitor.h b/src/arith/constr_visitor.h new file mode 100644 index 000000000000..779d34e4e988 --- /dev/null +++ b/src/arith/constr_visitor.h @@ -0,0 +1,554 @@ +/* + * 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. + */ + +/*! + * \file constr_visitor.h + * \brief Save arithmetic premises at different points of an IR traversal. + */ +#ifndef TVM_ARITH_CONSTR_VISITOR_H_ +#define TVM_ARITH_CONSTR_VISITOR_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace tvm { +namespace arith { + +/*! + * \brief Validate the restricted arithmetic domain used by snapshot proofs. + * + * Use a closed set of scalar signed int32/int64 arithmetic and boolean operators. + * Casts must preserve the signed integer value; multiplication requires a constant + * factor. Calls, loads, division, and remainders are excluded. + * Signed overflow follows the existing TIR undefined semantics. + * Division and remainder premises can introduce congruences whose combination + * overflows the analyzer even when every input expression is small. Until that + * growth is modeled, reject them in predicates, bindings, ranges, and queries. + * + * Also limit symbolic expansion: even a defined execution can have coefficients + * that overflow the analyzer's int64 arithmetic when products are distributed. + * Track a conservative magnitude through expressions and Bind definitions. Add + * magnitudes for sums/comparisons and multiply them for products. Growth must + * fit each integer node's type, including + * int32 nodes below widening casts. This deliberately excludes some valid large + * indices rather than asking the analyzer to construct overflowing literals. + */ +class ConstraintExprValidator : public tirx::ExprFunctor { + public: + bool IsSupported(const PrimExpr& expr) { + expression_bounds_.clear(); + return VisitExpr(expr) < kUnsupported; + } + + bool Bind(const tirx::Var& var, const PrimExpr& value) { + expression_bounds_.clear(); + uint64_t bound = CheckType(var->ty.as_or_throw(), VisitExpr(value)); + if (bound == kUnsupported) return false; + bindings_[var.get()] = bound; + return true; + } + + bool Bind(const tirx::Var& var, const Range& range) { + expression_bounds_.clear(); + uint64_t bound = CheckType(var->ty.as_or_throw(), + Add(VisitExpr(range->min), VisitExpr(range->extent))); + if (bound == kUnsupported) return false; + // A singleton range can become a value binding. Retain its dependencies + // even when the extent is only known to be one after simplification. + bindings_[var.get()] = bound; + return true; + } + + private: + static constexpr uint64_t kUnsupported = uint64_t{1} << 63; + + static uint64_t CheckType(PrimType ty, uint64_t bound) { + if (ty.MatchesCode(kDLInt) && bound >= (uint64_t{1} << (ty.bits() - 1))) { + return kUnsupported; + } + return bound; + } + + static uint64_t Add(uint64_t a, uint64_t b) { + return a >= kUnsupported - b ? kUnsupported : a + b; + } + + static uint64_t Mul(uint64_t a, uint64_t b) { + return a > (kUnsupported - 1) / b ? kUnsupported : a * b; + } + + uint64_t VisitExpr(const Expr& expr) final { + auto it = expression_bounds_.find(expr.get()); + if (it != expression_bounds_.end()) return it->second; + auto value = expr.as(); + if (!value) return kUnsupported; + PrimType ty = value.value().ty(); + if (!ty.IsScalar() || !(ty.MatchesCode(kDLBool) || + (ty.MatchesCode(kDLInt) && (ty.bits() == 32 || ty.bits() == 64)))) { + return kUnsupported; + } + if (ty.MatchesCode(kDLBool) && + !(expr.as() || expr.as() || expr.as() || + expr.as() || expr.as() || expr.as() || + expr.as() || expr.as() || expr.as() || + expr.as() || expr.as())) { + return kUnsupported; + } + uint64_t bound = CheckType(ty, ExprFunctor::VisitExpr(expr)); + expression_bounds_.emplace(expr.get(), bound); + return bound; + } + + uint64_t VisitExprDefault_(const ffi::Object*) final { return kUnsupported; } + + uint64_t VisitExpr_(const tirx::VarNode* op) final { + auto it = bindings_.find(op); + return it == bindings_.end() ? 1 : it->second; + } + + uint64_t VisitExpr_(const IntImmNode* op) final { + if (op->value == std::numeric_limits::min()) return kUnsupported; + return std::max(1, op->value < 0 ? -op->value : op->value); + } + + uint64_t VisitExpr_(const prim::CastNode* op) final { + PrimType from = op->value.ty(); + PrimType to = op->ty.as_or_throw(); + if (!(from.MatchesCode(kDLInt) && to.MatchesCode(kDLInt) && from.bits() <= to.bits())) { + return kUnsupported; + } + return VisitExpr(op->value); + } + +#define TVM_CONSTR_ADDITIVE_BOUND(Node) \ + uint64_t VisitExpr_(const prim::Node* op) final { \ + return Add(VisitExpr(op->a), VisitExpr(op->b)); \ + } + TVM_CONSTR_ADDITIVE_BOUND(AddNode) + TVM_CONSTR_ADDITIVE_BOUND(SubNode) + TVM_CONSTR_ADDITIVE_BOUND(MinNode) + TVM_CONSTR_ADDITIVE_BOUND(MaxNode) + TVM_CONSTR_ADDITIVE_BOUND(AndNode) + TVM_CONSTR_ADDITIVE_BOUND(OrNode) +#undef TVM_CONSTR_ADDITIVE_BOUND + + // Comparisons may be rewritten as a difference of their integer operands. + // The result is boolean, but the difference must fit the operands' type. +#define TVM_CONSTR_COMPARISON_BOUND(Node) \ + uint64_t VisitExpr_(const prim::Node* op) final { \ + return CheckType(op->a.ty(), Add(VisitExpr(op->a), VisitExpr(op->b))); \ + } + TVM_CONSTR_COMPARISON_BOUND(EQNode) + TVM_CONSTR_COMPARISON_BOUND(NENode) + TVM_CONSTR_COMPARISON_BOUND(LTNode) + TVM_CONSTR_COMPARISON_BOUND(LENode) + TVM_CONSTR_COMPARISON_BOUND(GTNode) + TVM_CONSTR_COMPARISON_BOUND(GENode) +#undef TVM_CONSTR_COMPARISON_BOUND + + uint64_t VisitExpr_(const prim::NotNode* op) final { return VisitExpr(op->a); } + + uint64_t VisitExpr_(const prim::MulNode* op) final { + if (!(op->a.as() || op->b.as())) return kUnsupported; + return Mul(VisitExpr(op->a), VisitExpr(op->b)); + } + + std::unordered_map bindings_; + // Memoize within one check to avoid expanding shared expression DAGs. + // Reset between checks, since bindings and expression lifetimes can change. + std::unordered_map expression_bounds_; +}; + +inline bool IsSupportedConstraintExpr(const PrimExpr& expr) { + return ConstraintExprValidator().IsSupported(expr); +} + +/*! \brief One premise, retaining bindings for the analyzer's rewrite and bound tables. */ +struct Constr { + enum Kind { kPredicate, kBindValue, kBindRange }; + + explicit Constr(PrimExpr predicate) : kind(kPredicate), value(std::move(predicate)) {} + Constr(tirx::Var var, PrimExpr value) + : kind(kBindValue), var(std::move(var)), value(std::move(value)) {} + Constr(tirx::Var var, Range range) + : kind(kBindRange), var(std::move(var)), range(std::move(range)) {} + + Kind kind; + tirx::Var var; + PrimExpr value{ffi::UnsafeInit{}}; + Range range; +}; + +/*! + * \brief An ordered snapshot of pure scalar premises. + * + * Unlike IRVisitorWithAnalyzer's current analyzer context, a snapshot can outlive + * the scope in which it was collected. Consumers must distinguish the variables + * of different executions before combining snapshots. Merge means conjunction + * of premises, not a control-flow join. + */ +struct ConstrSet { + std::vector constraints; + + ConstrSet RenameVars(const std::function& rename) const { + auto substitute = [&](const tirx::Var& var) -> ffi::Optional { + return rename(var).as_or_throw(); + }; + auto substitute_expr = [&](const PrimExpr& expr) { + return tirx::SubstituteWithDataTypeLegalization(expr, substitute); + }; + ConstrSet result; + for (const auto& c : constraints) { + switch (c.kind) { + case Constr::kPredicate: + result.constraints.emplace_back(substitute_expr(c.value)); + break; + case Constr::kBindValue: + result.constraints.emplace_back(rename(c.var), substitute_expr(c.value)); + break; + case Constr::kBindRange: + result.constraints.emplace_back( + rename(c.var), Range::FromMinExtent(substitute_expr(c.range->min), + substitute_expr(c.range->extent), c.range->span)); + break; + } + } + return result; + } + + ConstrSet Merge(const ConstrSet& other) const { + ConstrSet result = *this; + result.constraints.insert(result.constraints.end(), other.constraints.begin(), + other.constraints.end()); + return result; + } + + bool CanProve(const PrimExpr& predicate) const { + ConstraintExprValidator validator; + // Rebinding a variable can discard information or introduce inconsistent + // premises. Require the consumer to rename independent executions first. + std::unordered_set bound; + for (const auto& c : constraints) { + if (c.kind != Constr::kPredicate && !bound.insert(c.var.get()).second) { + return false; + } + // Validate replay as well as collection: a consumer can construct or + // rename snapshots without going through ConstrVisitor. + if (c.kind != Constr::kPredicate) { + auto var = c.var.as(); + if (!var || !IsSupportedConstraintExpr(var.value())) return false; + } + if (c.kind == Constr::kBindRange) { + if (!validator.Bind(c.var, c.range)) return false; + } else if (c.kind == Constr::kBindValue) { + if (!validator.Bind(c.var, c.value)) return false; + } else if (!validator.IsSupported(c.value)) { + return false; + } + } + if (!validator.IsSupported(predicate)) return false; + + // Inline immutable Bind values before replaying the snapshot. Analyzer::Bind + // installs a rewrite rule for the bound variable. Rewriting a comparison + // through several such rules can apply algebraic cancellation to an + // expression whose value is still symbolic, so it must not be used as the + // proof boundary. A value binding is an SSA definition; substituting its + // already-inlined value preserves that definition without asking the + // analyzer to rewrite through it. + std::unordered_map substitutions; + auto substitute = [&](const tirx::Var& var) -> ffi::Optional { + auto it = substitutions.find(var.get()); + if (it == substitutions.end()) return std::nullopt; + return it->second; + }; + auto inline_expr = [&](const PrimExpr& value) { + return substitutions.empty() ? value + : tirx::SubstituteWithDataTypeLegalization(value, substitute); + }; + + Analyzer analyzer; + // Congruence alone often separates flat addresses, e.g. even and odd + // indices. Try this inexpensive sufficient condition before replaying + // bindings and entering predicate scopes. No snapshot facts are assumed. + if (const auto* ne = predicate.as(); ne && ne->a.ty().MatchesCode(kDLInt)) { + if (analyzer->modular_set(ne->a - ne->b)->base != 0) return true; + } + WithGroup contexts; + for (const auto& c : constraints) { + switch (c.kind) { + case Constr::kPredicate: { + // Bound analysis needs normalized comparisons, e.g. !(x < n) -> x >= n. + PrimExpr value = inline_expr(c.value); + if (!IsSupportedConstraintExpr(value)) return false; + contexts.Emplace(analyzer, analyzer->rewrite_simplify(value)); + break; + } + case Constr::kBindValue: { + PrimExpr value = inline_expr(c.value); + if (!IsSupportedConstraintExpr(value)) return false; + substitutions.emplace(c.var.get(), value); + break; + } + case Constr::kBindRange: { + Range range = Range::FromMinExtent(inline_expr(c.range->min), + inline_expr(c.range->extent), c.range->span); + if (!IsSupportedConstraintExpr(range->min) || !IsSupportedConstraintExpr(range->extent)) { + return false; + } + if (tirx::is_one(range->extent)) { + // A singleton range is an SSA definition as well. Keep it on the + // same substitution path as BindValue so analyzer rewrite rules + // cannot cross a symbolic definition. + substitutions.emplace(c.var.get(), range->min); + } else { + analyzer->Bind(c.var, range); + } + break; + } + } + } + PrimExpr query = inline_expr(predicate); + if (!IsSupportedConstraintExpr(query)) return false; + return analyzer->CanProve(query); + } +}; + +/*! + * \brief Collect constraints that can safely be replayed at another access point. + * + * Only expressions accepted by IsSupportedConstraintExpr are retained. A Bind + * of a mutable read or an unsupported conversion does not create a persistent + * rewrite. Its variable remains an unconstrained symbol, which a consumer must + * rename per execution. Dropping such facts weakens the premises without + * introducing assumptions about memory snapshots or wrapping arithmetic. + * + * Derived visitors that override control flow must use WithConstrScope and add + * the appropriate premises after visiting conditions and bounds. + */ +class ConstrVisitor : public tirx::StmtExprVisitor { + public: + using StmtExprVisitor::VisitExpr_; + using StmtExprVisitor::VisitStmt_; + + ConstrSet GetConstrSet() const { return {constraints_}; } + + /*! + * \brief Save predicates and the bindings needed by an expression or predicate. + * + * Bindings are in definition order, so a backward pass also finds transitive + * dependencies. Keep every predicate, including ones after the definition of + * a query variable: they may constrain it indirectly through another variable. + * Unused bindings (often unrelated loop axes or scalar temporaries) do + * not need to be copied, renamed, validated, and replayed for an address proof. + * Omitting a premise only weakens the snapshot. + */ + ConstrSet GetConstrSet(const PrimExpr& expr) const { + std::unordered_set needed; + auto add_vars = [&](const PrimExpr& value) { + auto walk_fn = [&](const Var& var) -> ffi::Expected { + if (const auto* var_node = var.as()) needed.insert(var_node); + return ffi::WalkResult::Advance(); + }; + ffi::StructuralWalk(value, walk_fn); + }; + add_vars(expr); + for (const auto& c : constraints_) { + if (c.kind == Constr::kPredicate) add_vars(c.value); + } + ConstrSet result; + for (auto it = constraints_.rbegin(); it != constraints_.rend(); ++it) { + const auto& c = *it; + if (c.kind != Constr::kPredicate && !needed.count(c.var.get())) continue; + result.constraints.push_back(c); + if (c.kind == Constr::kBindRange) { + add_vars(c.range->min); + add_vars(c.range->extent); + } else if (c.kind == Constr::kBindValue) { + add_vars(c.value); + } + } + std::reverse(result.constraints.begin(), result.constraints.end()); + return result; + } + + void VisitStmt_(const tirx::BindNode* op) override { + this->VisitExpr(op->value); + AddBinding(op->var, op->value); + } + + void VisitStmt_(const tirx::AssertStmtNode* op) override { + StmtExprVisitor::VisitStmt_(op); + AddConstraint(op->condition); + } + + // SeqStmt does not introduce a scope. Bind/Assert facts also remain visible + // to following siblings when the sequence contains another SeqStmt. + + void VisitStmt_(const tirx::IfThenElseNode* op) override { + this->VisitExpr(op->condition); + WithConstrScope([&]() { + AddConstraint(op->condition); + this->VisitStmt(op->then_case); + }); + if (op->else_case) { + WithConstrScope([&]() { + AddConstraint(prim::Not(op->condition)); + this->VisitStmt(op->else_case.value()); + }); + } + } + + void VisitStmt_(const tirx::AttrStmtNode* op) override { + this->VisitExpr(op->value); + WithConstrScope([&]() { + if (op->attr_key == tirx::attr::thread_extent || + op->attr_key == s_tir::attr::virtual_thread) { + auto iv = op->node.as_or_throw(); + AddRange(iv->var, Range::FromMinExtent(IntImm(op->value.ty(), 0), op->value)); + } + this->VisitStmt(op->body); + }); + } + + void VisitStmt_(const tirx::ForNode* op) override { + this->VisitExpr(op->min); + this->VisitExpr(op->extent); + if (op->step) this->VisitExpr(op->step.value()); + WithConstrScope([&]() { + if (!op->step || tirx::is_one(op->step.value())) { + AddRange(op->loop_var, Range::FromMinExtent(op->min, op->extent)); + } + AddConstraint(op->extent > IntImm(op->extent.ty(), 0)); + this->VisitStmt(op->body); + }); + } + + void VisitStmt_(const tirx::WhileNode* op) override { + this->VisitExpr(op->condition); + WithConstrScope([&]() { + AddConstraint(op->condition); + this->VisitStmt(op->body); + }); + } + + void VisitStmt_(const tirx::SBlockNode* op) override { + WithConstrScope([&]() { + auto visit_region = [&](const tirx::BufferRegion& region) { + this->VisitBufferUse(region->buffer); + for (const auto& range : region->region) { + this->VisitExpr(range->min); + this->VisitExpr(range->extent); + } + }; + for (const auto& iv : op->iter_vars) { + this->VisitExpr(iv->dom->min); + this->VisitExpr(iv->dom->extent); + } + for (const auto& buffer : op->alloc_buffers) { + this->VisitBufferDef(buffer, /*alloc_data=*/true); + } + for (const auto& region : op->reads) visit_region(region); + for (const auto& region : op->writes) visit_region(region); + for (const auto& match : op->match_buffers) { + this->VisitBufferDef(match->buffer, /*alloc_data=*/true); + visit_region(match->source); + } + // Initialization executes only on the first reduction iteration. Its + // assertions and local bindings are not premises of subsequent updates. + if (op->init) { + WithConstrScope([&]() { this->VisitStmt(op->init.value()); }); + } + this->VisitStmt(op->body); + }); + } + + void VisitExpr_(const prim::LetNode* op) override { + this->VisitExpr(op->value); + WithConstrScope([&]() { + AddBinding(op->var, op->value); + this->VisitExpr(op->body); + }); + } + + // Select may evaluate both operands, so neither gets a branch constraint. + void VisitExpr_(const CallNode* op) override { + if (op->op.same_as(prim::builtin::if_then_else())) { + auto condition = op->args[0].as_or_throw(); + this->VisitExpr(condition); + WithConstrScope([&]() { + AddConstraint(condition); + this->VisitExpr(op->args[1]); + }); + WithConstrScope([&]() { + AddConstraint(prim::Not(condition)); + this->VisitExpr(op->args[2]); + }); + } else { + StmtExprVisitor::VisitExpr_(op); + } + } + + protected: + template + void WithConstrScope(F&& body) { + struct Guard { + std::vector& constraints; + size_t size; + ~Guard() { constraints.erase(constraints.begin() + size, constraints.end()); } + } guard{constraints_, constraints_.size()}; + body(); + } + + void AddConstraint(const PrimExpr& predicate) { + if (IsSupportedConstraintExpr(predicate)) constraints_.emplace_back(predicate); + } + + void AddBinding(const tirx::Var& var, const Expr& value) { + if (auto scalar = value.as(); scalar && IsSupportedConstraintExpr(scalar.value())) { + constraints_.emplace_back(var, scalar.value()); + } + } + + void AddRange(const tirx::Var& var, const Range& range) { + if (IsSupportedConstraintExpr(range->min) && IsSupportedConstraintExpr(range->extent)) { + constraints_.emplace_back(var, range); + } + } + + private: + std::vector constraints_; +}; + +} // namespace arith +} // namespace tvm +#endif // TVM_ARITH_CONSTR_VISITOR_H_ diff --git a/tests/cpp/arith_constr_visitor_test.cc b/tests/cpp/arith_constr_visitor_test.cc new file mode 100644 index 000000000000..86436baea83f --- /dev/null +++ b/tests/cpp/arith_constr_visitor_test.cc @@ -0,0 +1,291 @@ +/* + * 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 + +#include +#include + +#include "../../src/arith/constr_visitor.h" + +namespace tvm { +namespace arith { +namespace { +using namespace tirx; + +class SnapshotCollector : public ConstrVisitor { + public: + using ConstrVisitor::VisitExpr_; + using ConstrVisitor::VisitStmt_; + + void VisitStmt_(const EvaluateNode* op) override { + snapshots.push_back(GetConstrSet(op->value.as_or_throw())); + ConstrVisitor::VisitStmt_(op); + } + + void VisitExpr_(const TensorLoadNode* op) override { + loads.push_back(GetConstrSet(op->indices[0])); + ConstrVisitor::VisitExpr_(op); + } + + std::vector snapshots; + std::vector loads; +}; + +TEST(ConstrVisitor, SnapshotsOutliveBranchAndVisitor) { + PrimVar x("x", PrimType::Int(32)); + std::vector snapshots; + { + SnapshotCollector visitor; + visitor(SeqStmt({IfThenElse(x < 8, Evaluate(x), Evaluate(x)), Evaluate(x)})); + snapshots = visitor.snapshots; + } + ASSERT_EQ(snapshots.size(), 3); + EXPECT_TRUE(snapshots[0].CanProve(x < 8)); + EXPECT_TRUE(snapshots[1].CanProve(x >= 8)); + EXPECT_FALSE(snapshots[2].CanProve(x < 8)); + EXPECT_FALSE(snapshots[2].CanProve(x >= 8)); +} + +TEST(ConstrVisitor, LaterBindingsDoNotChangeEarlierSnapshots) { + PrimVar x("x", PrimType::Int(32)); + PrimVar y("y", PrimType::Int(32)); + SnapshotCollector visitor; + visitor(SeqStmt({Bind(x, PrimExpr(3)), Evaluate(x), Bind(y, x + 1), Evaluate(y)})); + ASSERT_EQ(visitor.snapshots.size(), 2); + EXPECT_TRUE(visitor.snapshots[0].CanProve(x == 3)); + EXPECT_FALSE(visitor.snapshots[0].CanProve(y == 4)); + EXPECT_TRUE(visitor.snapshots[1].CanProve(y == 4)); +} + +TEST(ConstrVisitor, AddressSnapshotKeepsTransitiveBindings) { + PrimVar x("x", PrimType::Int(32)); + PrimVar y("y", PrimType::Int(32)); + PrimVar z("z", PrimType::Int(32)); + ffi::Array statements{Bind(x, PrimExpr(3)), Bind(y, x + 1), Bind(z, y * 2)}; + for (int i = 0; i < 128; ++i) { + statements.push_back(Bind(PrimVar("unused", PrimType::Int(32)), x + i)); + } + statements.push_back(Evaluate(z)); + SnapshotCollector visitor; + visitor(SeqStmt(statements)); + ASSERT_EQ(visitor.snapshots.size(), 1); + EXPECT_EQ(visitor.snapshots[0].constraints.size(), 3); + EXPECT_TRUE(visitor.snapshots[0].CanProve(z == 8)); + EXPECT_EQ(visitor.GetConstrSet().constraints.size(), 131); +} + +TEST(ConstrVisitor, AddressSnapshotKeepsIndirectPredicateDependencies) { + PrimVar x("x", PrimType::Int(32)); + PrimVar query("query", PrimType::Int(32)); + PrimVar y("y", PrimType::Int(32)); + SnapshotCollector visitor; + visitor(SeqStmt({Bind(query, x + 1), Bind(y, x * 2), IfThenElse(y < 16, Evaluate(query))})); + ASSERT_EQ(visitor.snapshots.size(), 1); + EXPECT_TRUE(visitor.snapshots[0].CanProve(query < 9)); +} + +TEST(ConstrVisitor, RangeDependenciesSurviveLoopExit) { + PrimVar n("n", PrimType::Int(32)); + PrimVar i("i", PrimType::Int(32)); + SnapshotCollector visitor; + visitor(SeqStmt({Bind(n, PrimExpr(4)), For(i, 0, n, ForKind::kSerial, Evaluate(i))})); + ASSERT_EQ(visitor.snapshots.size(), 1); + EXPECT_TRUE(visitor.snapshots[0].CanProve(i >= 0 && i < 4)); + EXPECT_FALSE(visitor.GetConstrSet().CanProve(i < 4)); +} + +TEST(ConstrVisitor, LetBindingsAreScopedToTheBody) { + PrimVar x("x", PrimType::Int(32)); + auto buffer = decl_buffer({16}, PrimType::Int(32)); + SnapshotCollector visitor; + visitor(Evaluate(prim::Let(x, 3, BufferLoad(buffer, {x})))); + visitor(Evaluate(BufferLoad(buffer, {0}))); + ASSERT_EQ(visitor.loads.size(), 2); + EXPECT_TRUE(visitor.loads[0].CanProve(x == 3)); + EXPECT_FALSE(visitor.loads[1].CanProve(x == 3)); + EXPECT_FALSE(visitor.GetConstrSet().CanProve(x == 3)); +} + +TEST(ConstrVisitor, MutableReadDoesNotBecomeAPersistentBinding) { + PrimVar x("x", PrimType::Int(32)); + auto buffer = decl_buffer({16}, PrimType::Int(32)); + SnapshotCollector visitor; + visitor(SeqStmt({Bind(x, BufferLoad(buffer, {0})), Evaluate(x)})); + EXPECT_TRUE(visitor.GetConstrSet().constraints.empty()); + EXPECT_FALSE(visitor.snapshots[0].CanProve(x == 0)); +} + +TEST(ConstrSet, RenameDistinguishesBindingsAcrossSnapshots) { + PrimVar index("index", PrimType::Int(32)); + PrimVar mapped("mapped", PrimType::Int(32)); + ConstrSet first{{Constr(index, Range::FromMinExtent(0, 32)), Constr(mapped, 2 * index)}}; + std::unordered_map vars; + auto rename = [&](const Var& var) { + auto [it, inserted] = vars.emplace(var.get(), var); + if (inserted) it->second = var.CopyWithSuffix("_other"); + return it->second; + }; + auto second = first.RenameVars(rename); + auto other_mapped = rename(mapped).as_or_throw(); + auto merged = first.Merge(second); + EXPECT_TRUE(merged.CanProve(mapped != other_mapped + 1)); + EXPECT_FALSE(merged.CanProve(mapped != other_mapped)); + EXPECT_FALSE(merged.CanProve(mapped == other_mapped)); + EXPECT_TRUE(first.CanProve(mapped < 64)); +} + +TEST(ConstrSet, MergeDoesNotSilentlyDiscardDuplicateBindings) { + PrimVar x("x", PrimType::Int(32)); + ConstrSet first{{Constr(x, PrimExpr(0))}}; + ConstrSet second{{Constr(x, PrimExpr(1))}}; + EXPECT_FALSE(first.Merge(second).CanProve(x == 0)); + EXPECT_FALSE(first.Merge(first).CanProve(x == 0)); +} + +TEST(ConstrSet, ManuallyConstructedUnsupportedFactsAreRejected) { + PrimVar x("x", PrimType::Int(32)); + PrimVar y("y", PrimType::Int(32)); + PrimExpr wrapped = prim::Cast(PrimType::Int(32), prim::Cast(PrimType::UInt(8), x + 128)); + ConstrSet bindings{{Constr(y, wrapped)}}; + ConstrSet predicates{{Constr(wrapped < 128)}}; + ConstrSet ranges{{Constr(y, Range::FromMinExtent(0, wrapped))}}; + EXPECT_FALSE(bindings.CanProve(y == y)); + EXPECT_FALSE(predicates.CanProve(x == x)); + EXPECT_FALSE(ranges.CanProve(y == y)); + EXPECT_FALSE(ConstrSet{}.CanProve(wrapped == y)); +} + +TEST(ConstrSet, SymbolicGrowthAcrossBindingsIsRejected) { + PrimVar x("x", PrimType::Int(64)); + PrimVar j("j", PrimType::Int(64)); + PrimVar k("k", PrimType::Int(64)); + PrimExpr scale = IntImm::Int64(5000000000); + PrimExpr inner = scale * x - IntImm::Int64(9999999999); + // x == 2 gives inner == 1, and scale * inner == scale. All runtime + // intermediates fit int64, but distributing the products needs 25 * 10^18. + EXPECT_FALSE(ConstrSet{}.CanProve(scale * inner != scale)); + ConstrSet bindings{{Constr(j, inner), Constr(k, scale * j)}}; + EXPECT_FALSE(bindings.CanProve(k != scale)); + EXPECT_FALSE(bindings.CanProve(k - scale == 0)); +} + +TEST(ConstrSet, InlineBindingsBeforeReplay) { + PrimVar x("x", PrimType::Int(64)); + PrimVar w("w", PrimType::Int(64)); + PrimVar z("z", PrimType::Int(64)); + PrimVar y("y", PrimType::Int(64)); + PrimExpr scale = IntImm::Int64(1500000000); + PrimExpr value = IntImm::Int64(361500000000); + + // x == -8 and w == 7 satisfy y == value and y >= value. Replaying the + // Bind values through Analyzer::Bind used to simplify that predicate to + // w <= -241, incorrectly proving y != value. Inline SSA values before + // entering analyzer scopes so the snapshot remains a sound sufficient proof. + ConstrSet bindings{ + {Constr(z, x * IntImm::Int64(-31)), Constr(y, (z - w) * scale), Constr(y >= value)}}; + EXPECT_FALSE(bindings.CanProve(y != value)); + + ConstrSet singleton_range{{Constr(z, Range::FromMinExtent(x * IntImm::Int64(-31), 1)), + Constr(y, (z - w) * scale), Constr(y >= value)}}; + EXPECT_FALSE(singleton_range.CanProve(y != value)); +} + +TEST(ConstrSet, CombinedRemainderPremisesDoNotExcludeAValidValue) { + PrimVar x("x", PrimType::Int(64)); + PrimVar j("j", PrimType::Int(64)); + PrimExpr one = IntImm::Int64(1); + for (bool truncate : {false, true}) { + for (int64_t modulus : {3037000500LL, 4000000000LL}) { + auto remainder = [&](PrimExpr divisor) { + return truncate ? truncmod(x, divisor) : floormod(x, divisor); + }; + PrimExpr first = remainder(IntImm::Int64(modulus)); + PrimExpr second = remainder(IntImm::Int64(modulus + 1)); + // x == 1 satisfies both premises without any runtime overflow. Their + // combined modulus overflows int64 inside the analyzer's intersection. + ConstrSet facts{{Constr(first == one), Constr(second == one)}}; + EXPECT_FALSE(facts.CanProve(x != one)); + EXPECT_FALSE(facts.CanProve(IntImm::Int64(0) != x - one)); + // A remainder hidden in a Bind must not bypass replay validation. + ConstrSet bindings{{Constr(j, first), Constr(j == one), Constr(second == one)}}; + EXPECT_FALSE(bindings.CanProve(x != one)); + } + } +} + +TEST(ConstrVisitor, DroppingRemainderFactsPreservesIndependentPremises) { + PrimVar x("x", PrimType::Int(64)); + PrimVar tx("tx", PrimType::Int(32)); + PrimExpr one = IntImm::Int64(1); + SnapshotCollector visitor; + visitor(IfThenElse( + tx < 32, IfThenElse(floormod(x, IntImm::Int64(4000000000)) == one, + IfThenElse(floormod(x, IntImm::Int64(4000000001)) == one, Evaluate(x))))); + ASSERT_EQ(visitor.snapshots.size(), 1); + EXPECT_FALSE(visitor.snapshots[0].CanProve(x != one)); + EXPECT_TRUE(visitor.snapshots[0].CanProve(tx < 32)); +} + +TEST(ConstrVisitor, ReductionInitFactsDoNotReachTheUpdate) { + PrimVar i("i", PrimType::Int(32)); + SBlock block({IterVar(Range::FromMinExtent(0, 4), i, IterVarType::kCommReduce)}, {}, {}, "reduce", + Evaluate(i), + SeqStmt({AssertStmt(i == 0, prim::StringImm("AssertionError"), {}), Evaluate(i)})); + SnapshotCollector visitor; + visitor(block); + ASSERT_EQ(visitor.snapshots.size(), 2); + EXPECT_TRUE(visitor.snapshots[0].CanProve(i == 0)); + EXPECT_FALSE(visitor.snapshots[1].CanProve(i == 0)); + EXPECT_FALSE(visitor.GetConstrSet().CanProve(i == 0)); +} + +TEST(ConstrSet, ExpansionMustFitEachIntegerNode) { + PrimVar x("x", PrimType::Int(32)); + PrimVar j("j", PrimType::Int(32)); + PrimVar k("k", PrimType::Int(32)); + PrimExpr inner = 50000 * x - 99999; + PrimExpr expr = 50000 * inner - 50000; + EXPECT_FALSE(IsSupportedConstraintExpr(expr)); + EXPECT_FALSE(IsSupportedConstraintExpr(prim::Cast(PrimType::Int(64), expr))); + // Boolean comparisons can introduce a difference in their operands' dtype. + EXPECT_TRUE(IsSupportedConstraintExpr(1500000000 * x)); + EXPECT_FALSE(IsSupportedConstraintExpr(1500000000 * x != -1500000000 * x)); + ConstrSet bindings{{Constr(j, inner), Constr(k, 50000 * j)}}; + EXPECT_FALSE(bindings.CanProve(k != 50000)); + ConstrSet ranges{{Constr(j, inner), Constr(k, Range::FromMinExtent(50000 * j, 1))}}; + EXPECT_FALSE(ranges.CanProve(k != 50000)); + PrimExpr wide = IntImm::Int64(50000) * + (IntImm::Int64(50000) * prim::Cast(PrimType::Int(64), x) - IntImm::Int64(99999)); + EXPECT_TRUE(IsSupportedConstraintExpr(wide)); +} + +TEST(ConstrVisitor, SharedExpressionGrowthIsBounded) { + PrimVar x("x", PrimType::Int(64)); + PrimExpr expr = x; + for (int i = 0; i < 80; ++i) { + expr = prim::Add(expr, expr); + } + // The IR is small, but expanding it as a tree has 2^80 leaves. + EXPECT_FALSE(IsSupportedConstraintExpr(expr)); +} + +} // namespace +} // namespace arith +} // namespace tvm From f9036633d640acab2acd86e7626b488e46843c6d Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Tue, 15 Sep 2026 20:55:54 -0400 Subject: [PATCH 2/3] [Tests][Frontend] Trim redundant PyTorch frontend coverage --- .../test_frontend_from_exported_program.py | 1121 +++-------------- tests/python/relax/test_frontend_from_fx.py | 734 ++--------- 2 files changed, 290 insertions(+), 1565 deletions(-) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index c74573c49f14..cc58a02eba9e 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -62,35 +62,38 @@ def verify_model( tvm.ir.assert_structural_equal(mod, expected, map_free_vars=map_free_vars) -def verify_model_numerically(torch_model, example_args, rtol=1e-7, atol=1e-7): - """Verify model by comparing numerical outputs between PyTorch and TVM.""" - with torch.no_grad(): - pytorch_output = torch_model(*example_args) - - exported_program = export(torch_model, args=example_args) - mod = from_exported_program(exported_program) - target = tvm.target.Target("llvm") - ex = relax.build(mod, target) - vm = relax.VirtualMachine(ex, tvm.cpu()) - - tvm_args = [tvm.runtime.tensor(arg.numpy()) for arg in example_args] - tvm_output = vm["main"](*tvm_args) - - if hasattr(tvm_output, "numpy"): - tvm_output_np = tvm_output.numpy() - else: - tvm_output_np = tvm_output[0].numpy() - - pytorch_output_np = ( - pytorch_output.numpy() - if isinstance(pytorch_output, torch.Tensor) - else pytorch_output[0].numpy() - ) +def verify_model_numerically( + torch_model, + example_args, + rtol=1e-7, + atol=1e-7, + *, + dynamic_shapes=None, + input_sets=None, + run_ep_decomposition=True, +): + """Build once and compare every output, including runs with different input shapes.""" + if not env.has_llvm(): + pytest.skip("need llvm") + exported_program = export(torch_model, args=example_args, dynamic_shapes=dynamic_shapes) + mod = from_exported_program(exported_program, run_ep_decomposition=run_ep_decomposition) + vm = relax.VirtualMachine(relax.build(mod, target="llvm"), tvm.cpu()) - assert pytorch_output_np.shape == tvm_output_np.shape, ( - f"Shape mismatch: PyTorch {pytorch_output_np.shape} vs TVM {tvm_output_np.shape}" - ) - tvm.testing.assert_allclose(pytorch_output_np, tvm_output_np, rtol=rtol, atol=atol) + for args in (example_args,) if input_sets is None else input_sets: + tvm_args = [tvm.runtime.tensor(arg.detach().numpy()) for arg in args] + with torch.no_grad(): + expected = torch_model(*(arg.clone() for arg in args)) + expected = torch.utils._pytree.tree_leaves(expected) + actual = vm["main"](*tvm_args) + actual = [actual] if isinstance(actual, tvm.runtime.Tensor) else list(actual) + assert len(actual) == len(expected) + for actual_value, expected_value in zip(actual, expected): + actual_array = actual_value.numpy() + expected_array = expected_value.numpy() + assert actual_array.shape == expected_array.shape + assert actual_array.dtype == expected_array.dtype + np.testing.assert_allclose(actual_array, expected_array, rtol=rtol, atol=atol) + return mod operator_basic_unary = [ @@ -148,96 +151,47 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( verify_model(UnaryOp(), example_args, {}, expected) -def test_round_decimals(): - """torch.round(x, decimals) is exported as aten.round.decimals, which was missing - from the convert map (only round.default was registered) and made any explicit - decimals -- including decimals=0 -- fail with - "AssertionError: Unsupported function types ['round.decimals']". - - With the decimals overload registered, torch.round(x, decimals) must convert and - match PyTorch's round-half-to-even results, including negative decimals - (round(25, -1) == 20) where the scale-by-0.1 float precision path used to be wrong. - """ - - class RoundDecimalsModel(Module): - def __init__(self, decimals): - super().__init__() - self.decimals = decimals - - def forward(self, input): - return torch.round(input, decimals=self.decimals) - - # Half values exercise ties-to-even; 25/125/165 exercise the negative-decimals path. - x = torch.tensor( - [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25], dtype=torch.float32 - ) - for decimals in (0, 1, -1, -2): - verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,), rtol=1e-6, atol=1e-6) - - -def test_round_decimals_low_precision(): - """Scaling for low-precision inputs must happen in float32 and be cast back. - - 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 == 250000 - exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5 - already overflows float16 (the scale itself becomes inf), turning decimals=5 - and -5 into NaN. Upcasting the input to float32 keeps the scaling exact; the - rounded result is cast back to the input dtype. - """ +@pytest.mark.parametrize( + "dtype, decimals", + [ + pytest.param(torch.float32, (0, 1, -1), id="ties-to-even"), + pytest.param(torch.float16, (4, 5, -5), id="float16-scaling"), + ], +) +def test_round_decimals(dtype, decimals): + class RoundDecimals(Module): + def forward(self, x): + return tuple(torch.round(x, decimals=d) for d in decimals) - class RoundDecimalsModel(Module): - def __init__(self, decimals): - super().__init__() - self.decimals = decimals + # Negative decimals catch reciprocal-scaling errors; float16 catches overflow + # both in the scaled input (4) and in the scale itself (5 and -5). + x = torch.tensor([0.5, 1.5, 2.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25], dtype=dtype) + verify_model_numerically(RoundDecimals(), (x,), rtol=1e-6, atol=1e-6) - def forward(self, input): - return torch.round(input, decimals=self.decimals) - x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5], dtype=torch.float16) - # Positive decimals exercise the multiply-by-10**d overflow (4, 5); - # negative decimals exercise the 10**|d| scale overflowing float16 (-5). - for decimals in (2, 4, 5, -2, -4, -5): - verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,), rtol=1e-6, atol=1e-6) +def test_round_decimals_large(): + """An overflowing scale must remain importable in either direction.""" + class RoundDecimals(Module): + def forward(self, x): + return torch.round(x, decimals=309), torch.round(x, decimals=-309) -def test_round_decimals_large(): - """A large |decimals| must import and run without OverflowError. - - The scale 10**|decimals| used to be built as an unbounded host Python int - before being handed to relax.const, whose int-to-float conversion raises - OverflowError ("int too large to convert to float") once |decimals| >= 309 - (10**309 already exceeds the float64 range). PyTorch accepts such decimals and - exports a valid aten.round.decimals node, so importing the exported program - must not crash on them. The scale is now built directly in the float dtype and - saturates to inf once it leaves the finite range, matching PyTorch, whose - all-NaN result here comes from the same inf scale. - """ - - class RoundDecimalsModel(Module): - def __init__(self, decimals): - super().__init__() - self.decimals = decimals + @I.ir_module + class Expected: + @R.function + def main(x: R.Tensor((2,), "float32")): + with R.dataflow(): + scaled_up = R.multiply(x, R.const(float("inf"), "float32")) + rounded_up = R.round(scaled_up) + positive = R.divide(rounded_up, R.const(float("inf"), "float32")) + scaled_down = R.divide(x, R.const(float("inf"), "float32")) + rounded_down = R.round(scaled_down) + negative = R.multiply(rounded_down, R.const(float("inf"), "float32")) + result = (positive, negative) + R.output(result) + return result - def forward(self, input): - return torch.round(input, decimals=self.decimals) - - x = torch.tensor([0.5, 1.5, 25.0, -0.5, 0.0], dtype=torch.float32) - for decimals in (309, -309): - exported_program = export(RoundDecimalsModel(decimals).eval(), args=(x,)) - mod = from_exported_program(exported_program) # used to raise OverflowError here - ex = relax.build(mod, target="llvm") - vm = relax.VirtualMachine(ex, tvm.cpu()) - tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) - got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() - - # The scale overflows to inf, and IEEE arithmetic turns every element into - # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining - # (empty here) finite elements separately, since allclose fails on NaN. - expected = torch.round(x, decimals=decimals) - actual = torch.as_tensor(got) - assert torch.equal(torch.isnan(actual), torch.isnan(expected)) - finite = ~torch.isnan(expected) - assert torch.allclose(actual[finite], expected[finite], rtol=1e-6, atol=1e-6) + verify_model(RoundDecimals(), (torch.ones(2),), {}, Expected) operator_bool_unary = [ @@ -1052,10 +1006,18 @@ def main( verify_model(Atan2(), example_args, {}, expected) -def test_logical_and(): - class LogicalAnd(Module): +@pytest.mark.parametrize( + "torch_op, relax_op", + [ + (torch.logical_and, R.logical_and), + (torch.logical_or, R.logical_or), + (torch.logical_xor, R.logical_xor), + ], +) +def test_logical_binary(torch_op, relax_op): + class LogicalBinary(Module): def forward(self, lhs, rhs): - return torch.logical_and(lhs, rhs) + return torch_op(lhs, rhs) @tvm.script.ir_module class expected: @@ -1068,7 +1030,7 @@ def main( with R.dataflow(): lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") - lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_and(lv, lv1) + lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(lv, lv1) gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv2,) R.output(gv) return gv @@ -1077,7 +1039,7 @@ def main( torch.randn(1, 3, 10, 10, dtype=torch.float32), torch.randn(1, 3, 10, 10, dtype=torch.float32), ) - verify_model(LogicalAnd(), example_args, {}, expected) + verify_model(LogicalBinary(), example_args, {}, expected) def test_logical_not(): @@ -1103,62 +1065,6 @@ def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( verify_model(LogicalNot(), example_args, {}, expected) -def test_logical_or(): - class LogicalOr(Module): - def forward(self, lhs, rhs): - return torch.logical_or(lhs, rhs) - - @tvm.script.ir_module - class expected: - @R.function - def main( - lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), - rhs: R.Tensor((1, 3, 10, 10), dtype="float32"), - ) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") - lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") - lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_or(lv, lv1) - gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv2,) - R.output(gv) - return gv - - example_args = ( - torch.randn(1, 3, 10, 10, dtype=torch.float32), - torch.randn(1, 3, 10, 10, dtype=torch.float32), - ) - verify_model(LogicalOr(), example_args, {}, expected) - - -def test_logical_xor(): - class LogicalXor(Module): - def forward(self, lhs, rhs): - return torch.logical_xor(lhs, rhs) - - @tvm.script.ir_module - class expected: - @R.function - def main( - lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), - rhs: R.Tensor((1, 3, 10, 10), dtype="float32"), - ) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") - lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") - lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_xor(lv, lv1) - gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv2,) - R.output(gv) - return gv - - example_args = ( - torch.randn(1, 3, 10, 10, dtype=torch.float32), - torch.randn(1, 3, 10, 10, dtype=torch.float32), - ) - verify_model(LogicalXor(), example_args, {}, expected) - - def test_pow_integer(): class Pow(Module): def forward(self, input): @@ -2602,15 +2508,6 @@ def __init__(self): def forward(self, input): return self.conv(input) - class ConvTranspose1d1Func(Module): - def __init__(self): - super().__init__() - self.weight = torch.randn(size=[6, 6, 3]) - self.bias = torch.randn(size=[6]) - - def forward(self, input): - return torch.nn.functional.conv_transpose1d(input, self.weight, self.bias) - @tvm.script.ir_module class expected1: @R.function @@ -2678,10 +2575,6 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) - model = ConvTranspose1d1Func() - binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} - verify_model(model, example_args, binding, expected1) - model = ConvTranspose1d2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -2696,15 +2589,6 @@ def __init__(self): def forward(self, input): return self.conv(input) - class ConvTranspose2d1Func(Module): - def __init__(self): - super().__init__() - self.weight = torch.randn(size=[3, 3, 7, 7]) - self.bias = torch.randn(size=[3]) - - def forward(self, input): - return torch.nn.functional.conv_transpose2d(input, self.weight, self.bias) - @tvm.script.ir_module class expected1: @R.function @@ -2772,10 +2656,6 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) - model = ConvTranspose2d1Func() - binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} - verify_model(model, example_args, binding, expected1) - model = ConvTranspose2d2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -2790,15 +2670,6 @@ def __init__(self): def forward(self, input): return self.conv(input) - class Conv1D1Func(Module): - def __init__(self): - super().__init__() - self.weight = torch.randn(size=[6, 3, 7]) - self.bias = torch.randn(size=[6]) - - def forward(self, input): - return torch.nn.functional.conv1d(input, self.weight, self.bias) - @tvm.script.ir_module class expected1: @R.function @@ -2864,10 +2735,6 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) - model = Conv1D1Func() - binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} - verify_model(model, example_args, binding, expected1) - model = Conv1D2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -2882,15 +2749,6 @@ def __init__(self): def forward(self, input): return self.conv(input) - class Conv2D1Func(Module): - def __init__(self): - super().__init__() - self.weight = torch.randn(size=[6, 3, 7, 7]) - self.bias = torch.randn(size=[6]) - - def forward(self, input): - return torch.nn.functional.conv2d(input, self.weight, self.bias) - @tvm.script.ir_module class expected1: @R.function @@ -2956,10 +2814,6 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) - model = Conv2D1Func() - binding = {"w1": model.weight.numpy(), "w2": model.bias.numpy()} - verify_model(model, example_args, binding, expected1) - model = Conv2D2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -2974,15 +2828,6 @@ def __init__(self): def forward(self, input): return self.conv(input) - class Conv3D1Func(Module): - def __init__(self): - super().__init__() - self.weight = torch.randn(size=[6, 3, 7, 7, 7]) - self.bias = torch.randn(size=[6]) - - def forward(self, input): - return torch.nn.functional.conv3d(input, self.weight, self.bias) - @tvm.script.ir_module class expected1: @R.function @@ -3048,10 +2893,6 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) - model = Conv3D1Func() - binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} - verify_model(model, example_args, binding, expected1) - model = Conv3D2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -3420,19 +3261,7 @@ def main( def test_einsum_repeated_subscript(): - """einsum with repeated subscripts (diagonal / trace) on the default - decomposition path. - - ``run_decompositions`` (default) lowers repeated-subscript einsum to - ``aten.diagonal`` + ``permute`` (+ ``sum`` for the trace), which the - frontend converts with the ``_diagonal`` lowering. For the zero-offset - square case (e.g. ``torch.einsum("ii->i")`` on an ``N x N`` input) the - frontend emits a single repeated-subscript einsum that reads the diagonal - directly; otherwise it permutes the diagonal dims to the trailing two axes, - slices each to the diagonal length, and runs an einsum ``...zz->...z``. - This used to raise ``AssertionError: Unsupported function types - ['diagonal.default']``. - """ + """Decomposed diagonal extraction must lower to a single einsum.""" class EinsumDiag(Module): def __init__(self): @@ -3456,118 +3285,20 @@ def main(x: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((3,), dtype=" example_args = (torch.randn(3, 3, dtype=torch.float32),) verify_model(EinsumDiag(), example_args, {}, Expected) - class TraceEinsum(Module): - def forward(self, x): - return torch.einsum("ii->", x) - - class BatchedDiagEinsum(Module): - def forward(self, x): - return torch.einsum("...ii->...i", x) - - class AttentionEinsum(Module): - def forward(self, x, y): - return torch.einsum("abca,abcb->c", x, y) - - verify_model_numerically(TraceEinsum(), (torch.randn(4, 4),)) - verify_model_numerically(BatchedDiagEinsum(), (torch.randn(2, 3, 3),)) - verify_model_numerically(AttentionEinsum(), (torch.randn(3, 3, 4, 3), torch.randn(3, 3, 4, 3))) - - class DirectDiagonal(Module): - def __init__(self): - super().__init__() - self.offset = 1 - + class BatchedDiagonal(Module): def forward(self, x): - return torch.diagonal(x, self.offset, 0, 1) + return torch.einsum("...ii->...i", x), torch.einsum("...ii->...", x) - class DirectTrace(Module): - def forward(self, x): - return torch.trace(x) - - verify_model_numerically(DirectDiagonal(), (torch.randn(3, 4),)) - verify_model_numerically(DirectTrace(), (torch.randn(4, 4),)) + verify_model_numerically(BatchedDiagonal(), (torch.arange(18.0).reshape(2, 3, 3),)) - # Out-of-range offsets (|offset| >= max(extent1, extent2)) are valid in - # PyTorch and yield an empty diagonal of shape (0,); the lowering must - # clamp the diagonal length to zero instead of producing negative slice - # extents or a wrong non-empty shape. - class DirectDiagonalOutOfRange(Module): - def __init__(self, offset): - super().__init__() - self.offset = offset +def test_diagonal_offsets(): + class Diagonal(Module): def forward(self, x): - return torch.diagonal(x, self.offset, 0, 1) + # Non-square input, both offset signs, and one empty result per sign. + return tuple(torch.diagonal(x, offset, 0, 1) for offset in (1, -1, 4, -3)) - for offset in [4, 5, 6, -3, -4, -5, -6]: - verify_model_numerically(DirectDiagonalOutOfRange(offset), (torch.randn(3, 4),)) - - -def test_einsum_diagonal_lowers_without_full_size_intermediate(): - """Regression test: a zero-offset square diagonal must not materialize - full-size intermediates. - - ``torch.einsum("ii->i")`` on an ``N x N`` input is decomposed to - ``aten.diagonal`` by ``run_decompositions``. Lowering that diagonal by - permuting the diagonal dims to the trailing axes, slicing each to the - diagonal length, and running the ``...zz->...z`` einsum materializes three - full-size ``N x N`` intermediates (an identity permute and two identity - strided slices) and hence three O(N^2) copy loops before the final O(N) - diagonal loop. The ``_diagonal`` fast path instead emits a single - repeated-subscript einsum that reads the diagonal directly, so no full-size - intermediate exists in the frontend graph (and therefore neither in the - lowered TIR). Assert that every intermediate produced by a call is at most - O(N), both before and after legalization. - """ - - class EinsumDiag(Module): - def forward(self, x): - return torch.einsum("ii->i", x) - - n = 8 - exported_program = export(EinsumDiag(), args=(torch.randn(n, n),)) - mod = from_exported_program(exported_program) - - def rank2_call_results(ir_mod): - """Names of calls whose result is a rank-2 (full-size) tensor.""" - results = [] - for func in ir_mod.functions.values(): - if not isinstance(func, relax.Function): - continue - for block in func.body.blocks: - for binding in block.bindings: - if not ( - isinstance(binding.value, relax.Call) - and isinstance(binding.value.op, tvm.ir.Op) - ): - continue - if isinstance(binding.var.ty, relax.TensorType) and binding.var.ty.ndim == 2: - results.append(binding.value.op.name) - return results - - # The diagonal must be the only full-size (N x N) tensor touched: it is the - # function input read directly by a single repeated-subscript einsum. No - # call may produce a rank-2 intermediate. - assert rank2_call_results(mod) == [] - - # Sanity check that the graph really performs the diagonal: exactly one - # einsum on the N x N input producing an N-vector. - einsum_calls = [] - for block in mod["main"].body.blocks: - for binding in block.bindings: - if ( - isinstance(binding.value, relax.Call) - and isinstance(binding.value.op, tvm.ir.Op) - and binding.value.op.name == "relax.einsum" - ): - einsum_calls.append(binding.var) - assert len(einsum_calls) == 1 - assert einsum_calls[0].ty.ndim == 1 - - # Legalize and check again on the lowered graph. - with tvm.target.Target("llvm"): - lowered = relax.transform.LegalizeOps()(mod) - assert rank2_call_results(lowered) == [] + verify_model_numerically(Diagonal(), (torch.arange(12.0).reshape(3, 4),)) def test_outer(): @@ -3761,15 +3492,6 @@ def __init__(self): def forward(self, input): return self.linear(input) - class Dense1Func(Module): - def __init__(self): - super().__init__() - self.weight = torch.randn(size=[7, 10]) - self.bias = torch.randn(size=[7]) - - def forward(self, input): - return torch.nn.functional.linear(input, self.weight, self.bias) - @tvm.script.ir_module class expected1: @R.function @@ -3824,10 +3546,6 @@ def main( binding = {"w1": model.linear.weight.detach().numpy(), "w2": model.linear.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) - model = Dense1Func() - binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} - verify_model(model, example_args, binding, expected1) - model = Dense2() binding = {"w1": model.linear.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -3923,13 +3641,6 @@ def __init__(self): def forward(self, input): return self.pool(input) - class MaxPool2d_functional(Module): - def __init__(self): - super().__init__() - - def forward(self, input): - return torch.nn.functional.max_pool2d(input, kernel_size=[1, 1]) - @tvm.script.ir_module class expected1: @R.function @@ -4027,7 +3738,6 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),) verify_model(MaxPool2d(), example_args, {}, expected1) - verify_model(MaxPool2d_functional(), example_args, {}, expected1) verify_model(MaxPool2d2(), example_args, {}, expected2) verify_model(MaxPool2d3(), example_args, {}, expected3) @@ -4041,13 +3751,6 @@ def __init__(self): def forward(self, input): return self.pool(input) - class MaxPool3d_functional(Module): - def __init__(self): - super().__init__() - - def forward(self, input): - return torch.nn.functional.max_pool3d(input, kernel_size=[1, 1, 1]) - @tvm.script.ir_module class expected1: @R.function @@ -4149,7 +3852,6 @@ def main(input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32")) -> R.Tuple( # Verify the models with expected IR modules verify_model(MaxPool3d(), example_args1, {}, expected1) - verify_model(MaxPool3d_functional(), example_args1, {}, expected1) verify_model(MaxPool3d2(), example_args2, {}, expected2) verify_model(MaxPool3d3(), example_args3, {}, expected3) @@ -5816,37 +5518,25 @@ def forward(self, x): example_args = (torch.randn(1, 4, 3, dtype=torch.float32),) tokens = torch.export.Dim("tokens", min=1, max=8) - exported_program = export( + mod = verify_model_numerically( DynamicShapeOps(), - args=example_args, + example_args, dynamic_shapes={"x": {1: tokens}}, + input_sets=[(torch.randn(1, token_count, 3),) for token_count in (4, 6)], + rtol=0, + atol=0, ) - mod = from_exported_program(exported_program) - script = mod.script() assert "R.tensor_to_shape" in script assert "R.shape_to_tensor" in script - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - for token_count in (4, 6): - torch_input = torch.randn(1, token_count, 3) - expected = DynamicShapeOps()(torch_input) - actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) - for actual_value, expected_value in zip(actual, expected): - np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) - @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @pytest.mark.parametrize( ("item_dtype", "item_shape"), [ - (torch.int8, ()), - (torch.uint8, ()), - (torch.int16, ()), (torch.int32, ()), - (torch.int64, ()), - (torch.int64, (1,)), + (torch.uint8, (1,)), (torch.int64, (1, 1)), ], ) @@ -5867,21 +5557,15 @@ def forward(self, x): rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) example_args = (torch.randn(3, 4, dtype=torch.float32),) - exported_program = export( + verify_model_numerically( DynamicItem(), - args=example_args, + example_args, + rtol=0, + atol=0, dynamic_shapes={"x": {0: rows, 1: columns}}, + input_sets=[(torch.randn(shape),) for shape in ((3, 4), (5, 2))], + run_ep_decomposition=False, ) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - for shape in ((3, 4), (5, 2)): - torch_input = torch.randn(shape, dtype=torch.float32) - expected = DynamicItem()(torch_input) - actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) - for actual_value, expected_value in zip(actual, expected): - np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -5900,15 +5584,9 @@ def forward(self, x, signed, unsigned, wide): torch.tensor([200], dtype=torch.uint8), torch.tensor(1 << 40, dtype=torch.int64), ) - exported_program = export(RuntimeItems(), args=example_args) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - expected = RuntimeItems()(*example_args) - actual = vm["main"](*(tvm.runtime.tensor(arg.numpy()) for arg in example_args)) - for actual_value, expected_value in zip(actual, expected): - np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + verify_model_numerically( + RuntimeItems(), example_args, rtol=0, atol=0, run_ep_decomposition=False + ) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -5925,26 +5603,15 @@ def forward(self, x): example_args = (torch.randn(3, 4, dtype=torch.float32),) rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) - exported_program = export( + verify_model_numerically( DynamicFills(), - args=example_args, + example_args, + rtol=0, + atol=0, dynamic_shapes={"x": {0: rows, 1: columns}}, + input_sets=[(torch.randn(shape),) for shape in ((3, 4), (5, 2))], + run_ep_decomposition=False, ) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - for shape in ((3, 4), (5, 2)): - torch_input = torch.randn(shape, dtype=torch.float32) - expected = DynamicFills()(torch_input) - actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) - actual_arrays = [value.numpy() for value in actual] - - assert actual_arrays[0].dtype == np.dtype("int64") - assert actual_arrays[1].dtype == np.dtype("float64") - for actual_value, expected_value in zip(actual_arrays, expected): - assert actual_value.dtype == expected_value.numpy().dtype - np.testing.assert_array_equal(actual_value, expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -5968,21 +5635,15 @@ def forward(self, x): rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) example_args = (torch.randn(3, 4, dtype=torch.float32),) - exported_program = export( + verify_model_numerically( DynamicBooleanFills(), - args=example_args, + example_args, + rtol=0, + atol=0, dynamic_shapes={"x": {0: rows, 1: columns}}, + input_sets=[(torch.randn(shape),) for shape in ((3, 3), (3, 4))], + run_ep_decomposition=False, ) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - for shape in ((3, 3), (3, 4)): - torch_input = torch.randn(shape, dtype=torch.float32) - expected = DynamicBooleanFills()(torch_input) - actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) - for actual_value, expected_value in zip(actual, expected): - np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -6005,21 +5666,15 @@ def forward(self, x): rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) example_args = (torch.randn(3, 4, dtype=torch.float32),) - exported_program = export( + verify_model_numerically( DynamicScalarArithmetic(), - args=example_args, + example_args, + rtol=0, + atol=0, dynamic_shapes={"x": {0: rows, 1: columns}}, + input_sets=[(torch.randn(shape),) for shape in ((3, 4), (5, 2))], + run_ep_decomposition=False, ) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - for shape in ((3, 4), (5, 2)): - torch_input = torch.randn(shape, dtype=torch.float32) - expected = DynamicScalarArithmetic()(torch_input) - actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) - for actual_value, expected_value in zip(actual, expected): - np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -6039,16 +5694,7 @@ def forward(self, x): ) example_args = (torch.randn(2, 3, dtype=torch.float32),) - exported_program = export(FullLike(), args=example_args) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - expected = FullLike()(*example_args) - actual = vm["main"](tvm.runtime.tensor(example_args[0].numpy())) - for actual_value, expected_value in zip(actual, expected): - assert actual_value.numpy().dtype == expected_value.numpy().dtype - np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + verify_model_numerically(FullLike(), example_args, rtol=0, atol=0, run_ep_decomposition=False) def test_split(): @@ -6086,16 +5732,7 @@ def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( def test_split_int_split_size(): - """x.split(int, dim) must produce chunks of size `split_size` (the last one - smaller when the dimension is not divisible), matching PyTorch. - - The frontend used to convert the int per-chunk size into a section count and - pass it as relax.op.split's int argument, which means "split into N equal - sections"; that yields wrong chunk shapes whenever - ceil(D / ceil(D / split_size)) != split_size (e.g. split_size > D/2). The - int branch now builds cumulative cut positions, the same as the list/tuple - form. - """ + """A non-divisible split size denotes chunk length, not section count.""" class Split6(Module): def forward(self, input): @@ -6126,43 +5763,6 @@ def main(input: R.Tensor((10,), dtype="float32")) -> R.Tuple( verify_model(Split6(), example_args, {}, Expected) # Differential check against native PyTorch for non-divisible sizes and dims. - class SplitModel(Module): - def __init__(self, split_size, dim): - super().__init__() - self.split_size = split_size - self.dim = dim - - def forward(self, input): - return input.split(self.split_size, dim=self.dim) - - def run_tvm(model, args): - exported_program = export(model, args=args) - mod = from_exported_program(exported_program) - ex = relax.build(mod, target="llvm") - vm = relax.VirtualMachine(ex, tvm.cpu()) - out = vm["main"](*[tvm.runtime.tensor(a.numpy()) for a in args]) - if hasattr(out, "numpy"): - return [out.numpy()] - return [o.numpy() for o in out] - - for shape, split_size, dim in [ - ((10,), 6, 0), - ((10,), 7, 0), - ((10,), 8, 0), - ((10,), 9, 0), - ((12,), 7, 0), - ((12, 8), 5, 1), - ((3, 10), 6, -1), - ]: - x = torch.arange(1, int(np.prod(shape)) + 1, dtype=torch.float32).reshape(shape) - refs = [r.numpy() for r in x.split(split_size, dim)] - outs = run_tvm(SplitModel(split_size, dim), (x,)) - assert [r.shape for r in refs] == [o.shape for o in outs], ( - f"split shape={shape} s={split_size} dim={dim}: " - f"torch {[r.shape for r in refs]} vs tvm {[o.shape for o in outs]}" - ) - for r, o in zip(refs, outs): - tvm.testing.assert_allclose(o, r, rtol=1e-7, atol=1e-7) def test_squeeze(): @@ -6667,21 +6267,15 @@ def forward(self, x): rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) example_args = (torch.randn(3, 4, dtype=torch.float32),) - exported_program = export( + verify_model_numerically( DynamicMaskedFills(), - args=example_args, + example_args, + rtol=0, + atol=0, dynamic_shapes={"x": {0: rows, 1: columns}}, + input_sets=[(torch.randn(shape),) for shape in ((3, 4), (5, 2))], + run_ep_decomposition=False, ) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - for shape in ((3, 4), (5, 2)): - torch_input = torch.randn(shape, dtype=torch.float32) - expected = DynamicMaskedFills()(torch_input) - actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) - for actual_value, expected_value in zip(actual, expected): - np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -6707,16 +6301,9 @@ def forward(self, x, mask): torch.arange(6, dtype=dtype).reshape(2, 3), torch.tensor([[True, False, True], [False, True, False]]), ) - exported_program = export(MaskedFills(), args=example_args) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - expected = MaskedFills()(*example_args) - actual = vm["main"](*(tvm.runtime.tensor(arg.numpy()) for arg in example_args)) - for actual_value, expected_value in zip(actual, expected): - assert actual_value.numpy().dtype == expected_value.numpy().dtype - np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + verify_model_numerically( + MaskedFills(), example_args, rtol=0, atol=0, run_ep_decomposition=False + ) def test_masked_select(): @@ -7270,18 +6857,10 @@ class Gather0(Module): def forward(self, data, indices): return torch.gather(data, 0, indices) - class Gather1(Module): - def forward(self, data, indices): - return torch.gather(data, 1, indices) - class Gather2(Module): def forward(self, data, indices): return torch.gather(data, -1, indices) - class Gather3(Module): - def forward(self, data, indices): - return torch.gather(data, -2, indices) - @tvm.script.ir_module class Expected0: @R.function @@ -7295,19 +6874,6 @@ def main( R.output(gv) return gv - @tvm.script.ir_module - class Expected1: - @R.function - def main( - inp_0: R.Tensor((2, 3), dtype="float32"), - inp_1: R.Tensor((2, 3), dtype="int64"), - ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")): - with R.dataflow(): - lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=1) - gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,) - R.output(gv) - return gv - @tvm.script.ir_module class Expected2: @R.function @@ -7321,28 +6887,13 @@ def main( R.output(gv) return gv - @tvm.script.ir_module - class Expected3: - @R.function - def main( - inp_0: R.Tensor((2, 3), dtype="float32"), - inp_1: R.Tensor((2, 3), dtype="int64"), - ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")): - with R.dataflow(): - lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=-2) - gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,) - R.output(gv) - return gv - example_args = ( torch.randn(2, 3, dtype=torch.float32), torch.randint(0, 3, (2, 3), dtype=torch.int64), ) verify_model(Gather0(), example_args, {}, Expected0) - verify_model(Gather1(), example_args, {}, Expected1) verify_model(Gather2(), example_args, {}, Expected2) - verify_model(Gather3(), example_args, {}, Expected3) def test_index_put(): @@ -7405,112 +6956,10 @@ def main( return gv # Test case 3: 3D input - class IndexPut3D(Module): - def forward(self, data, indices_0, indices_1, indices_2, values): - indices_tuple = (indices_0, indices_1, indices_2) - return data.index_put_(indices_tuple, values, accumulate=False) - - example_args_3d = ( - torch.randn(16, 32, 64, dtype=torch.float32), - torch.randint(0, 16, (128,), dtype=torch.int64), - torch.randint(0, 32, (128,), dtype=torch.int64), - torch.randint(0, 64, (128,), dtype=torch.int64), - torch.randn(128, dtype=torch.float32), - ) - - @I.ir_module - class Expected3D: - @R.function - def main( - data: R.Tensor((16, 32, 64), dtype="float32"), - indices_0: R.Tensor((128,), dtype="int64"), - indices_1: R.Tensor((128,), dtype="int64"), - indices_2: R.Tensor((128,), dtype="int64"), - values: R.Tensor((128,), dtype="float32"), - ) -> R.Tuple(R.Tensor((16, 32, 64), dtype="float32")): - with R.dataflow(): - lv: R.Tensor((16, 32, 64), dtype="float32") = R.index_put( - data, (indices_0, indices_1, indices_2), values, accumulate=False - ) - gv: R.Tuple(R.Tensor((16, 32, 64), dtype="float32")) = (lv,) - R.output(gv) - return gv # Test case 4: 4D input - class IndexPut4D(Module): - def forward(self, data, indices_0, indices_1, indices_2, indices_3, values): - indices_tuple = (indices_0, indices_1, indices_2, indices_3) - return data.index_put_(indices_tuple, values, accumulate=False) - - example_args_4d = ( - torch.randn(8, 16, 32, 64, dtype=torch.float32), - torch.randint(0, 8, (128,), dtype=torch.int64), - torch.randint(0, 16, (128,), dtype=torch.int64), - torch.randint(0, 32, (128,), dtype=torch.int64), - torch.randint(0, 64, (128,), dtype=torch.int64), - torch.randn(128, dtype=torch.float32), - ) - - @I.ir_module - class Expected4D: - @R.function - def main( - data: R.Tensor((8, 16, 32, 64), dtype="float32"), - indices_0: R.Tensor((128,), dtype="int64"), - indices_1: R.Tensor((128,), dtype="int64"), - indices_2: R.Tensor((128,), dtype="int64"), - indices_3: R.Tensor((128,), dtype="int64"), - values: R.Tensor((128,), dtype="float32"), - ) -> R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")): - with R.dataflow(): - lv: R.Tensor((8, 16, 32, 64), dtype="float32") = R.index_put( - data, - (indices_0, indices_1, indices_2, indices_3), - values, - accumulate=False, - ) - gv: R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")) = (lv,) - R.output(gv) - return gv # Test case 5: 5D input - class IndexPut5D(Module): - def forward(self, data, indices_0, indices_1, indices_2, indices_3, indices_4, values): - indices_tuple = (indices_0, indices_1, indices_2, indices_3, indices_4) - return data.index_put_(indices_tuple, values, accumulate=False) - - example_args_5d = ( - torch.randn(4, 8, 16, 32, 64, dtype=torch.float32), - torch.randint(0, 4, (128,), dtype=torch.int64), - torch.randint(0, 8, (128,), dtype=torch.int64), - torch.randint(0, 16, (128,), dtype=torch.int64), - torch.randint(0, 32, (128,), dtype=torch.int64), - torch.randint(0, 64, (128,), dtype=torch.int64), - torch.randn(128, dtype=torch.float32), - ) - - @I.ir_module - class Expected5D: - @R.function - def main( - data: R.Tensor((4, 8, 16, 32, 64), dtype="float32"), - indices_0: R.Tensor((128,), dtype="int64"), - indices_1: R.Tensor((128,), dtype="int64"), - indices_2: R.Tensor((128,), dtype="int64"), - indices_3: R.Tensor((128,), dtype="int64"), - indices_4: R.Tensor((128,), dtype="int64"), - values: R.Tensor((128,), dtype="float32"), - ) -> R.Tuple(R.Tensor((4, 8, 16, 32, 64), dtype="float32")): - with R.dataflow(): - lv: R.Tensor((4, 8, 16, 32, 64), dtype="float32") = R.index_put( - data, - (indices_0, indices_1, indices_2, indices_3, indices_4), - values, - accumulate=False, - ) - gv: R.Tuple(R.Tensor((4, 8, 16, 32, 64), dtype="float32")) = (lv,) - R.output(gv) - return gv # Test case 6: 2D input with multi-dimensional index (broadcasting) # This tests the multi-dimensional index support with broadcasting @@ -7660,9 +7109,6 @@ def main(x: R.Tensor((2, 10), dtype="float32")) -> R.Tuple( # Run verification for each case verify_model(IndexPut1D(), example_args_1d, {}, Expected1D) verify_model(IndexPut2D(), example_args_2d, {}, Expected2D) - verify_model(IndexPut3D(), example_args_3d, {}, Expected3D) - verify_model(IndexPut4D(), example_args_4d, {}, Expected4D) - verify_model(IndexPut5D(), example_args_5d, {}, Expected5D) verify_model(IndexPutBroadcast1D(), example_args_broadcast1, {}, ExpectedBroadcast1D) verify_model(IndexPutBroadcast2D(), example_args_broadcast2, {}, ExpectedBroadcast2D) verify_model(IndexPutBroadcast3D(), example_args_broadcast3d, {}, ExpectedBroadcast3D) @@ -7696,40 +7142,6 @@ def forward(self, x, buf, idx): ) -def test_m4d_diag_index_put_tuple_output_regression(): - class M4D(Module): - def forward(self, x): - b, k, n = 2, 3, 5 - buf = x.new_zeros(b, k, n, n) - idx = torch.arange(n, device=x.device) - - diag = buf[..., idx, idx] - diag = torch.nn.functional.elu(diag) + 1.0 + 1e-8 - buf[..., idx, idx] = diag - - return x[..., :1], buf - - ex_in = torch.zeros(2, 3, 5, dtype=torch.float32) - exported_program = export(M4D().eval(), args=(ex_in,)) - - exported_targets = [str(getattr(n, "target", "")) for n in exported_program.graph.nodes] - assert any("index_put" in target for target in exported_targets) - - # Regression focus: importing this graph should not segfault at Tuple construction. - mod = from_exported_program(exported_program) - ret_ty = mod["main"].ret_ty - assert isinstance(ret_ty, relax.TupleType) - - tensor_fields = [f for f in ret_ty.fields if isinstance(f, relax.TensorType)] - assert len(tensor_fields) >= 2 - # x: (2, 3, 5) → x[..., :1]: (2, 3, 1) - assert any(len(f.shape) == 3 and int(f.shape[-1]) == 1 for f in tensor_fields) - # buf: (2, 3, 5, 5) → 4-D with spatial dims 5x5 - assert any( - len(f.shape) == 4 and int(f.shape[-2]) == 5 and int(f.shape[-1]) == 5 for f in tensor_fields - ) - - def test_index_put_mutation_through_alias_regression(): class IndexPutAlias(Module): def forward(self, x, idx, values): @@ -7778,10 +7190,6 @@ class Flip0(Module): def forward(self, data): return torch.flip(data, [0]) - class Flip1(Module): - def forward(self, data): - return torch.flip(data, [1]) - @tvm.script.ir_module class Expected0: @R.function @@ -7794,22 +7202,9 @@ def main( R.output(gv) return gv - @tvm.script.ir_module - class Expected1: - @R.function - def main( - inp_0: R.Tensor((2, 2), dtype="float32"), - ) -> R.Tuple(R.Tensor((2, 2), dtype="float32")): - with R.dataflow(): - lv: R.Tensor((2, 2), dtype="float32") = R.flip(inp_0, axis=1) - gv: R.Tuple(R.Tensor((2, 2), dtype="float32")) = (lv,) - R.output(gv) - return gv - example_args = (torch.randn(2, 2, dtype=torch.float32),) verify_model(Flip0(), example_args, {}, Expected0) - verify_model(Flip1(), example_args, {}, Expected1) def test_flip_multi_axis(): @@ -8086,8 +7481,7 @@ def main( verify_model(Bucketize(), (input_tensor, boundaries), {}, Expected) -@pytest.mark.parametrize("right", [False, True]) -@pytest.mark.parametrize("out_int32", [False, True]) +@pytest.mark.parametrize("right, out_int32", [(False, False), (True, True)]) def test_bucketize_numerically(right, out_int32): class Bucketize(Module): def forward(self, input_tensor, boundaries): @@ -8605,42 +7999,34 @@ def main( verify_model(SparseMatrixMultiply(), example_args, {}, Expected) -@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") -def test_lstm(): - class LSTM(nn.Module): - def __init__(self, input_size, hidden_size, batch_first, bidirectional): +@pytest.mark.parametrize("rnn_type", [nn.LSTM, nn.GRU, nn.RNN], ids=["lstm", "gru", "rnn-tanh"]) +@pytest.mark.parametrize( + "batch_first, bidirectional", + [(True, False), (False, True)], + ids=["batch-first", "bidirectional"], +) +def test_recurrent(rnn_type, batch_first, bidirectional): + class Recurrent(Module): + def __init__(self): super().__init__() - self.lstm = nn.LSTM( - input_size=input_size, - hidden_size=hidden_size, - num_layers=1, - batch_first=batch_first, - bidirectional=bidirectional, - ) + self.rnn = rnn_type(3, 4, batch_first=batch_first, bidirectional=bidirectional) def forward(self, x): - y, _ = self.lstm(x) - return y - - # Unidirectional LSTM with batch_first=True - torch.manual_seed(42) - x = torch.randn(2, 3, 4, dtype=torch.float32) - verify_model_numerically(LSTM(4, 8, batch_first=True, bidirectional=False), (x,)) - - # Unidirectional LSTM with batch_first=False - torch.manual_seed(43) - x2 = torch.randn(4, 2, 3, dtype=torch.float32) - verify_model_numerically(LSTM(3, 6, batch_first=False, bidirectional=False), (x2,)) - - # Bidirectional LSTM with batch_first=True - torch.manual_seed(44) - x3 = torch.randn(2, 3, 4, dtype=torch.float32) - verify_model_numerically(LSTM(4, 8, batch_first=True, bidirectional=True), (x3,)) - - # Bidirectional LSTM with batch_first=False - torch.manual_seed(45) - x4 = torch.randn(4, 2, 3, dtype=torch.float32) - verify_model_numerically(LSTM(3, 6, batch_first=False, bidirectional=True), (x4,)) + output, state = self.rnn(x) + return (output, state) if rnn_type is nn.RNN else output + + # Exercise both layouts and direction counts without repeating their product. + # Retain the RNN hidden-state check alongside its sequence output. + with torch.random.fork_rng(devices=[]): + torch.manual_seed(42) + x = torch.randn(2, 3, 3) if batch_first else torch.randn(3, 2, 3) + verify_model_numerically( + Recurrent(), + (x,), + rtol=1e-4, + atol=1e-5, + run_ep_decomposition=rnn_type is not nn.RNN, + ) def test_tensor_none_tuple(): @@ -8665,96 +8051,6 @@ def main(x: R.Tensor((3,), dtype="float32")) -> R.Tuple( verify_model(TensorNoneModel(), example_args, {}, Expected) -def test_gru(): - class GRU(nn.Module): - def __init__(self, input_size, hidden_size, batch_first, bidirectional): - super().__init__() - self.gru = nn.GRU( - input_size=input_size, - hidden_size=hidden_size, - num_layers=1, - batch_first=batch_first, - bidirectional=bidirectional, - ) - - def forward(self, x): - y, _ = self.gru(x) - return y - - cases = [ - (42, (2, 3, 4), 4, 8, True, False), - (43, (4, 2, 3), 3, 6, False, False), - (44, (2, 3, 4), 4, 5, True, True), - (45, (4, 2, 3), 3, 4, False, True), - ] - for seed, shape, input_size, hidden_size, batch_first, bidirectional in cases: - torch.manual_seed(seed) - x = torch.randn(*shape, dtype=torch.float32) - verify_model_numerically( - GRU(input_size, hidden_size, batch_first, bidirectional), - (x,), - rtol=1e-4, - atol=1e-5, - ) - - -@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") -def test_rnn_tanh(): - target = tvm.target.Target("llvm") - - def _check(rnn_kwargs, x_shape, seed): - class RNNWithState(nn.Module): - def __init__(self): - super().__init__() - self.rnn = nn.RNN(nonlinearity="tanh", num_layers=1, **rnn_kwargs) - - def forward(self, x): - output, h_n = self.rnn(x) - return output, h_n - - torch.manual_seed(seed) - x = torch.randn(*x_shape, dtype=torch.float32) - model = RNNWithState() - with torch.no_grad(): - pt_out, pt_hn = model(x) - - exported_program = export(model, args=(x,)) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - ex = relax.build(mod, target) - vm = relax.VirtualMachine(ex, tvm.cpu()) - tvm_outputs = vm["main"](tvm.runtime.tensor(x.numpy())) - tvm_out_np = tvm_outputs[0].numpy() - tvm_hn_np = tvm_outputs[1].numpy() - - assert pt_out.shape == tvm_out_np.shape, ( - f"output shape mismatch: PyTorch {tuple(pt_out.shape)} vs TVM {tvm_out_np.shape}" - ) - assert pt_hn.shape == tvm_hn_np.shape, ( - f"h_n shape mismatch: PyTorch {tuple(pt_hn.shape)} vs TVM {tvm_hn_np.shape}" - ) - tvm.testing.assert_allclose(pt_out.numpy(), tvm_out_np, rtol=1e-4, atol=1e-5) - tvm.testing.assert_allclose(pt_hn.numpy(), tvm_hn_np, rtol=1e-4, atol=1e-5) - - # batch_first, unidirectional - _check( - {"input_size": 4, "hidden_size": 8, "batch_first": True, "bidirectional": False}, - (2, 3, 4), - seed=42, - ) - # seq-first (batch_first=False), unidirectional - _check( - {"input_size": 3, "hidden_size": 6, "batch_first": False, "bidirectional": False}, - (4, 2, 3), - seed=43, - ) - # bidirectional, batch_first - _check( - {"input_size": 4, "hidden_size": 8, "batch_first": True, "bidirectional": True}, - (2, 3, 4), - seed=44, - ) - - def test_dynamic_shape_with_range_constraints(): class DynamicModel(torch.nn.Module): def forward(self, x1, x2): @@ -9212,10 +8508,6 @@ class UpsampleNearest2dScale(Module): def forward(self, input): return torch.nn.functional.interpolate(input, scale_factor=2.0, mode="nearest") - class UpsampleNearest2dSize(Module): - def forward(self, input): - return torch.nn.functional.interpolate(input, size=(20, 20), mode="nearest") - example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),) @tvm.script.ir_module @@ -9236,26 +8528,7 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( R.output(gv) return gv - @tvm.script.ir_module - class expected_size: - @R.function - def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( - R.Tensor((1, 3, 20, 20), dtype="float32") - ): - with R.dataflow(): - lv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.image.resize2d( - input_1, - size=(20, 20), - layout="NCHW", - method="nearest_neighbor", - coordinate_transformation_mode="half_pixel", - ) - gv: R.Tuple(R.Tensor((1, 3, 20, 20), dtype="float32")) = (lv,) - R.output(gv) - return gv - verify_model(UpsampleNearest2dScale(), example_args, {}, expected_scale) - verify_model(UpsampleNearest2dSize(), example_args, {}, expected_size) def test_from_exported_program_sparse_csr_buffer(): @@ -9392,38 +8665,9 @@ def main( ) -@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") -def test_cond_shape_equality_predicate(): - class CondShapeEqualityModel(Module): - def forward(self, x): - def true_fn(x): - return x + 1.0 - - def false_fn(x): - return x - 1.0 - - return torch.cond(x.shape[0] == x.shape[1], true_fn, false_fn, (x,)) - - rows = torch.export.Dim("rows", min=1, max=8) - columns = torch.export.Dim("columns", min=1, max=8) - exported_program = export( - CondShapeEqualityModel(), - args=(torch.zeros(3, 3),), - dynamic_shapes={"x": {0: rows, 1: columns}}, - ) - mod = from_exported_program(exported_program) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - for shape, expected_value in (((3, 3), 1.0), ((2, 3), -1.0)): - torch_input = torch.zeros(shape, dtype=torch.float32) - actual = vm["main"](tvm.runtime.tensor(torch_input.numpy()))[0] - np.testing.assert_array_equal(actual.numpy(), np.full(shape, expected_value, "float32")) - - -@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") -def test_cond_shape_inequality_predicate(): - class CondShapeInequalityModel(Module): +@pytest.mark.parametrize("compare", [operator.eq, operator.ne], ids=["equal", "not-equal"]) +def test_cond_shape_comparison(compare): + class CondShapeModel(Module): def forward(self, x): def true_fn(x): return x + 1.0 @@ -9431,23 +8675,18 @@ def true_fn(x): def false_fn(x): return x - 1.0 - return torch.cond(x.shape[0] != x.shape[1], true_fn, false_fn, (x,)) + return torch.cond(compare(x.shape[0], x.shape[1]), true_fn, false_fn, (x,)) rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) - exported_program = export( - CondShapeInequalityModel(), - args=(torch.zeros(2, 3),), + verify_model_numerically( + CondShapeModel(), + (torch.zeros(2, 3),), dynamic_shapes={"x": {0: rows, 1: columns}}, + input_sets=[(torch.zeros(shape),) for shape in ((2, 3), (3, 3))], + rtol=0, + atol=0, ) - mod = from_exported_program(exported_program) - executable = relax.build(mod, tvm.target.Target("llvm")) - vm = relax.VirtualMachine(executable, tvm.cpu()) - - for shape, expected_value in (((2, 3), 1.0), ((3, 3), -1.0)): - torch_input = torch.zeros(shape, dtype=torch.float32) - actual = vm["main"](tvm.runtime.tensor(torch_input.numpy()))[0] - np.testing.assert_array_equal(actual.numpy(), np.full(shape, expected_value, "float32")) def test_cond_tuple_output(): @@ -9588,23 +8827,13 @@ class AffineGrid(Module): def forward(self, theta): return torch.nn.functional.affine_grid(theta, [2, 3, 8, 12], align_corners=True) - model = AffineGrid() - example_args = (torch.randn(2, 2, 3, dtype=torch.float32),) - - with torch.no_grad(): - pytorch_output = model(*example_args) - - exported_program = export(model, args=example_args) - mod = from_exported_program(exported_program, run_ep_decomposition=False) - - exe = tvm.compile(mod, target="llvm") - vm = relax.VirtualMachine(exe, tvm.cpu()) - - tvm_args = [tvm.runtime.tensor(arg.numpy()) for arg in example_args] - tvm_output = vm["main"](*tvm_args) - tvm_output_np = tvm_output[0].numpy() - - tvm.testing.assert_allclose(tvm_output_np, pytorch_output.numpy(), rtol=1e-5, atol=1e-5) + verify_model_numerically( + AffineGrid(), + (torch.randn(2, 2, 3),), + rtol=1e-5, + atol=1e-5, + run_ep_decomposition=False, + ) if __name__ == "__main__": diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index a0ba7971db0a..2126032c4fcf 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -33,10 +33,10 @@ from tvm.script import tirx as T -def verify_model(torch_model, input_info, binding, expected): +def verify_model(torch_model, input_info, binding, expected, **import_options): graph_model = fx.symbolic_trace(torch_model) with torch.no_grad(): - mod = from_fx(graph_model, input_info) + mod = from_fx(graph_model, input_info, **import_options) binding = {k: tvm.runtime.tensor(v) for k, v in binding.items()} expected = relax.transform.BindParams("main", binding)(expected) tvm.ir.assert_structural_equal(mod, expected) @@ -2007,46 +2007,6 @@ def main( binding = {} verify_model(model, input_info, binding, expected2) - class LayerNorm3(Module): - def __init__(self, shape): - super().__init__() - self.shape = shape - self.weight = torch.nn.Parameter(torch.ones(shape)) - self.bias = torch.nn.Parameter(torch.zeros(shape)) - - def forward(self, input): - return torch.nn.functional.layer_norm(input, self.shape, self.weight, self.bias, 1e-5) - - @tvm.script.ir_module - class expected3: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32"), - w1: R.Tensor([10, 10], dtype="float32"), - w2: R.Tensor([10, 10], dtype="float32"), - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm( - input_1, - w1, - w2, - axes=[-2, -1], - epsilon=1e-05, - center=True, - scale=True, - ) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv - - model = LayerNorm3([10, 10]) - binding = { - "w1": model.weight.detach().numpy(), - "w2": model.bias.detach().numpy(), - } - verify_model(model, input_info, binding, expected3) - def test_cross_entropy(): input_info = [([3, 2], "float32"), ([3], "int32")] @@ -2556,112 +2516,6 @@ def main( verify_model(DivFloorModel(), input_info, {}, expected_div_floor) -def test_round_decimals(): - """torch.round(x, decimals) through from_fx must match PyTorch's round-half-to-even - results, including negative decimals (round(25, -1) == 20). The previous - scale-by-10**decimals implementation multiplied by 0.1 for negative decimals, which - is numerically wrong: 25 * 0.1 == 2.5000000000000004 in float64 rounds up to 30. - """ - input_info = [([10], "float32")] - x = torch.tensor( - [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25], dtype=torch.float32 - ) - - class RoundDecimalsModel(Module): - def __init__(self, decimals): - super().__init__() - self.decimals = decimals - - def forward(self, input): - return torch.round(input, decimals=self.decimals) - - for decimals in (0, 1, -1, -2): - gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval()) - mod = from_fx(gm, input_info) - ex = relax.build(mod, target="llvm") - vm = relax.VirtualMachine(ex, tvm.cpu()) - tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) - got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() - tvm.testing.assert_allclose( - got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6, atol=1e-6 - ) - - -def test_round_decimals_low_precision(): - """Scaling for low-precision inputs must happen in float32 and be cast back. - - 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 == 250000 - exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5 - already overflows float16 (the scale itself becomes inf), turning decimals=5 - and -5 into NaN. Upcasting the input to float32 keeps the scaling exact; the - rounded result is cast back to the input dtype. - """ - input_info = [([8], "float16")] - x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5], dtype=torch.float16) - - class RoundDecimalsModel(Module): - def __init__(self, decimals): - super().__init__() - self.decimals = decimals - - def forward(self, input): - return torch.round(input, decimals=self.decimals) - - # Positive decimals exercise the multiply-by-10**d overflow (4, 5); - # negative decimals exercise the 10**|d| scale overflowing float16 (-5). - for decimals in (2, 4, 5, -2, -4, -5): - gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval()) - mod = from_fx(gm, input_info) - ex = relax.build(mod, target="llvm") - vm = relax.VirtualMachine(ex, tvm.cpu()) - tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) - got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() - tvm.testing.assert_allclose( - got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6, atol=1e-6 - ) - - -def test_round_decimals_large(): - """A large |decimals| must import and run without OverflowError. - - The scale 10**|decimals| used to be built as an unbounded host Python int - before being handed to relax.const, whose int-to-float conversion raises - OverflowError ("int too large to convert to float") once |decimals| >= 309 - (10**309 already exceeds the float64 range). PyTorch accepts such decimals -- - torch.round(x, decimals=309) -- and traces a valid round.decimals call, so - importing the graph must not crash on them. The scale is now built directly - in the float dtype and saturates to inf once it leaves the finite range, - matching PyTorch, whose all-NaN result here comes from the same inf scale. - """ - input_info = [([5], "float32")] - x = torch.tensor([0.5, 1.5, 25.0, -0.5, 0.0], dtype=torch.float32) - - class RoundDecimalsModel(Module): - def __init__(self, decimals): - super().__init__() - self.decimals = decimals - - def forward(self, input): - return torch.round(input, decimals=self.decimals) - - for decimals in (309, -309): - gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval()) - mod = from_fx(gm, input_info) # used to raise OverflowError here - ex = relax.build(mod, target="llvm") - vm = relax.VirtualMachine(ex, tvm.cpu()) - tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) - got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() - - # The scale overflows to inf, and IEEE arithmetic turns every element into - # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining - # (empty here) finite elements separately, since allclose fails on NaN. - expected = torch.round(x, decimals=decimals) - actual = torch.as_tensor(got) - assert torch.equal(torch.isnan(actual), torch.isnan(expected)) - finite = ~torch.isnan(expected) - assert torch.allclose(actual[finite], expected[finite], rtol=1e-6, atol=1e-6) - - def test_size(): input_info = [([1, 3, 10, 10], "float32")] @@ -3646,62 +3500,20 @@ def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( verify_model(Trunc(), input_info, {}, expected_trunc) -def test_logical_and(): - input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")] - - class LogicalAnd(Module): - def forward(self, lhs, rhs): - return torch.logical_and(lhs, rhs) - - @tvm.script.ir_module - class expected: - @R.function - def main( - lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), - rhs: R.Tensor((1, 3, 10, 10), dtype="float32"), - ) -> R.Tensor((1, 3, 10, 10), dtype="bool"): - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") - lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") - lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_and(lv, lv1) - gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2 - R.output(gv) - return gv - - verify_model(LogicalAnd(), input_info, {}, expected) - - -def test_logical_or(): - input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")] - - class LogicalOr(Module): - def forward(self, lhs, rhs): - return torch.logical_or(lhs, rhs) - - @tvm.script.ir_module - class expected: - @R.function - def main( - lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), - rhs: R.Tensor((1, 3, 10, 10), dtype="float32"), - ) -> R.Tensor((1, 3, 10, 10), dtype="bool"): - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") - lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") - lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_or(lv, lv1) - gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2 - R.output(gv) - return gv - - verify_model(LogicalOr(), input_info, {}, expected) - - -def test_logical_xor(): +@pytest.mark.parametrize( + "torch_op, relax_op", + [ + (torch.logical_and, R.logical_and), + (torch.logical_or, R.logical_or), + (torch.logical_xor, R.logical_xor), + ], +) +def test_logical_binary(torch_op, relax_op): input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")] - class LogicalXor(Module): + class LogicalBinary(Module): def forward(self, lhs, rhs): - return torch.logical_xor(lhs, rhs) + return torch_op(lhs, rhs) @tvm.script.ir_module class expected: @@ -3713,12 +3525,12 @@ def main( with R.dataflow(): lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") - lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_xor(lv, lv1) + lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(lv, lv1) gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2 R.output(gv) return gv - verify_model(LogicalXor(), input_info, {}, expected) + verify_model(LogicalBinary(), input_info, {}, expected) def test_pow_integer(): @@ -3743,403 +3555,100 @@ def main(inp_0: R.Tensor((4,), dtype="int64")) -> R.Tensor((4,), dtype="int64"): verify_model(Pow(), input_info, {}, expected) -def test_interpolate(): - input_info = [([1, 3, 10, 10], "float32")] - +@pytest.mark.parametrize( + "shape, layout, kwargs, size, method, coordinate_mode", + [ + ((1, 3, 10, 10), "NCHW", {"size": (5, 5)}, (5, 5), "nearest_neighbor", "asymmetric"), + ( + (1, 3, 10, 10), + "NCHW", + {"scale_factor": 2.0, "mode": "bilinear", "align_corners": False}, + (20, 20), + "linear", + "half_pixel", + ), + ( + (1, 3, 10, 10), + "NCHW", + {"scale_factor": (2.0, 1.0), "mode": "bicubic", "align_corners": False}, + (20, 10), + "cubic", + "half_pixel", + ), + ( + (1, 3, 4, 10, 10), + "NCDHW", + {"scale_factor": (2.0, 4.0, 4.0), "mode": "trilinear", "align_corners": False}, + (8, 40, 40), + "linear", + "half_pixel", + ), + ( + (1, 3, 4, 10, 10), + "NCDHW", + {"size": (8, 40, 40), "mode": "trilinear", "align_corners": True}, + (8, 40, 40), + "linear", + "align_corners", + ), + ((1, 10, 10, 3), "NHWC", {"size": (5, 5)}, (5, 5), "nearest_neighbor", "asymmetric"), + ( + (1, 10, 10, 3), + "NHWC", + {"scale_factor": 2.0, "mode": "bilinear", "align_corners": False}, + (20, 20), + "linear", + "half_pixel", + ), + ( + (1, 4, 10, 10, 3), + "NDHWC", + {"scale_factor": (2.0, 4.0, 4.0), "mode": "trilinear", "align_corners": True}, + (8, 40, 40), + "linear", + "align_corners", + ), + ], + ids=[ + "nearest", + "scalar-scale", + "tuple-scale-cubic", + "trilinear-scale", + "trilinear-size-aligned", + "nhwc-size", + "nhwc-scale", + "ndhwc-scale-aligned", + ], +) +def test_interpolate(shape, layout, kwargs, size, method, coordinate_mode): class Interpolate(Module): - def forward(self, input): - return torch.nn.functional.interpolate(input, (5, 5)) - - @tvm.script.ir_module - class expected1: - @R.function - def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( - (1, 3, 5, 5), dtype="float32" - ): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 5, 5), dtype="float32") = R.image.resize2d( - input_1, - (5, 5), - roi=[0.000000, 0.000000, 0.000000, 0.000000], - layout="NCHW", - method="nearest_neighbor", - coordinate_transformation_mode="asymmetric", - rounding_method="round", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 3, 5, 5), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Interpolate(), input_info, {}, expected1) - - class Interpolate2(Module): - def forward(self, input): - return torch.nn.functional.interpolate( - input, - size=None, - scale_factor=2.0, - mode="bilinear", - align_corners=False, - ) - - @tvm.script.ir_module - class expected2: - @R.function - def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( - (1, 3, 20, 20), dtype="float32" - ): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.image.resize2d( - input_1, - (20, 20), - roi=[0.000000, 0.000000, 0.000000, 0.000000], - layout="NCHW", - method="linear", - coordinate_transformation_mode="half_pixel", - rounding_method="round", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 3, 20, 20), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Interpolate2(), input_info, {}, expected2) - - class Interpolate3(Module): - def forward(self, input): - return torch.nn.functional.interpolate( - input, - size=None, - scale_factor=(2.0, 1.0), - mode="bilinear", - align_corners=False, - ) - - @tvm.script.ir_module - class expected3: - @R.function - def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( - (1, 3, 20, 10), dtype="float32" - ): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 20, 10), dtype="float32") = R.image.resize2d( - input_1, - (20, 10), - roi=[0.000000, 0.000000, 0.000000, 0.000000], - layout="NCHW", - method="linear", - coordinate_transformation_mode="half_pixel", - rounding_method="round", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 3, 20, 10), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Interpolate3(), input_info, {}, expected3) - - class Interpolate4(Module): - def forward(self, input): - return torch.nn.functional.interpolate( - input, - size=None, - scale_factor=(2.0, 1.0), - mode="bicubic", - align_corners=False, - ) - - @tvm.script.ir_module - class expected4: - @R.function - def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( - (1, 3, 20, 10), dtype="float32" - ): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 20, 10), dtype="float32") = R.image.resize2d( - input_1, - (20, 10), - roi=[0.000000, 0.000000, 0.000000, 0.000000], - layout="NCHW", - method="cubic", - coordinate_transformation_mode="half_pixel", - rounding_method="round", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 3, 20, 10), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Interpolate4(), input_info, {}, expected4) - - input_info_5d = [([1, 3, 4, 10, 10], "float32")] - - class Interpolate6(Module): - def forward(self, input): - return torch.nn.functional.interpolate( - input, - size=None, - scale_factor=(2.0, 4.0, 4.0), - mode="trilinear", - align_corners=False, - ) - - @tvm.script.ir_module - class expected6: - @R.function - def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor( - (1, 3, 8, 40, 40), dtype="float32" - ): - with R.dataflow(): - lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d( - input_5, - (8, 40, 40), - roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], - layout="NCDHW", - method="linear", - coordinate_transformation_mode="half_pixel", - rounding_method="", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Interpolate6(), input_info_5d, {}, expected6) - - class Interpolate7(Module): - def forward(self, input): - return torch.nn.functional.interpolate( - input, - size=(8, 40, 40), - mode="trilinear", - align_corners=False, - ) - - @tvm.script.ir_module - class expected7: - @R.function - def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor( - (1, 3, 8, 40, 40), dtype="float32" - ): - with R.dataflow(): - lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d( - input_5, - (8, 40, 40), - roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], - layout="NCDHW", - method="linear", - coordinate_transformation_mode="half_pixel", - rounding_method="", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Interpolate7(), input_info_5d, {}, expected7) - - class Interpolate8(Module): - def forward(self, input): - return torch.nn.functional.interpolate( - input, - size=(8, 40, 40), - mode="trilinear", - align_corners=True, - ) - - @tvm.script.ir_module - class expected8: - @R.function - def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor( - (1, 3, 8, 40, 40), dtype="float32" - ): - with R.dataflow(): - lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d( - input_5, - (8, 40, 40), - roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], - layout="NCDHW", - method="linear", - coordinate_transformation_mode="align_corners", - rounding_method="", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Interpolate8(), input_info_5d, {}, expected8) - - -def test_interpolate_nhwc_layout(): - input_info = [([1, 10, 10, 3], "float32")] - - class InterpolateNHWC(Module): - def forward(self, input): - return torch.nn.functional.interpolate(input, (5, 5)) - - @tvm.script.ir_module - class expected_nhwc: - @R.function - def main(input_1: R.Tensor((1, 10, 10, 3), dtype="float32")) -> R.Tensor( - (1, 5, 5, 3), dtype="float32" - ): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 5, 5, 3), dtype="float32") = R.image.resize2d( - input_1, - (5, 5), - roi=[0.000000, 0.000000, 0.000000, 0.000000], - layout="NHWC", - method="nearest_neighbor", - coordinate_transformation_mode="asymmetric", - rounding_method="round", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 5, 5, 3), dtype="float32") = lv - R.output(gv) - return gv - - # Test with NHWC layout - graph_model = fx.symbolic_trace(InterpolateNHWC()) - with torch.no_grad(): - mod = from_fx(graph_model, input_info, default_image_layout="NHWC") - tvm.ir.assert_structural_equal(mod, expected_nhwc) - - # Test with bilinear interpolation and NHWC layout - class InterpolateNHWC2(Module): - def forward(self, input): - return torch.nn.functional.interpolate( - input, size=None, scale_factor=2.0, mode="bilinear", align_corners=False - ) - - @tvm.script.ir_module - class expected_nhwc2: - @R.function - def main(input_1: R.Tensor((1, 10, 10, 3), dtype="float32")) -> R.Tensor( - (1, 20, 20, 3), dtype="float32" - ): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 20, 20, 3), dtype="float32") = R.image.resize2d( - input_1, - (20, 20), - roi=[0.000000, 0.000000, 0.000000, 0.000000], - layout="NHWC", - method="linear", - coordinate_transformation_mode="half_pixel", - rounding_method="round", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 20, 20, 3), dtype="float32") = lv - R.output(gv) - return gv - - graph_model2 = fx.symbolic_trace(InterpolateNHWC2()) - with torch.no_grad(): - mod2 = from_fx(graph_model2, input_info, default_image_layout="NHWC") - tvm.ir.assert_structural_equal(mod2, expected_nhwc2) - - input_info_5d = [([1, 4, 10, 10, 3], "float32")] + def forward(self, x): + return F.interpolate(x, **kwargs) - class InterpolateNHWC3(Module): - def forward(self, input): - return torch.nn.functional.interpolate( - input, - size=None, - scale_factor=(2.0, 4.0, 4.0), - mode="trilinear", - align_corners=False, - ) + resize = R.image.resize3d if len(shape) == 5 else R.image.resize2d + rounding_method = "" if len(shape) == 5 else "round" - @tvm.script.ir_module - class expected_nhwc3: - @R.function - def main(input_5: R.Tensor((1, 4, 10, 10, 3), dtype="float32")) -> R.Tensor( - (1, 8, 40, 40, 3), dtype="float32" - ): - with R.dataflow(): - lv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = R.image.resize3d( - input_5, - (8, 40, 40), - roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], - layout="NDHWC", - method="linear", - coordinate_transformation_mode="half_pixel", - rounding_method="", + x = relax.Var("x", relax.TensorType(shape, "float32")) + builder = relax.BlockBuilder() + with builder.function("main", [x]): + with builder.dataflow(): + resized = builder.emit( + resize( + x, + size, + layout=layout, + method=method, + coordinate_transformation_mode=coordinate_mode, + rounding_method=rounding_method, cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, ) - gv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = lv - R.output(gv) - return gv - - graph_model3 = fx.symbolic_trace(InterpolateNHWC3()) - with torch.no_grad(): - mod3 = from_fx(graph_model3, input_info_5d, default_image_layout="NDHWC") - tvm.ir.assert_structural_equal(mod3, expected_nhwc3) - - class InterpolateNHWC4(Module): - def forward(self, input): - return torch.nn.functional.interpolate( - input, - size=None, - scale_factor=(2.0, 4.0, 4.0), - mode="trilinear", - align_corners=True, ) + output = builder.emit_output(resized) + builder.emit_func_output(output) - @tvm.script.ir_module - class expected_nhwc4: - @R.function - def main(input_5: R.Tensor((1, 4, 10, 10, 3), dtype="float32")) -> R.Tensor( - (1, 8, 40, 40, 3), dtype="float32" - ): - with R.dataflow(): - lv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = R.image.resize3d( - input_5, - (8, 40, 40), - roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], - layout="NDHWC", - method="linear", - coordinate_transformation_mode="align_corners", - rounding_method="", - cubic_alpha=-0.75, - cubic_exclude=0, - extrapolation_value=0, - ) - gv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = lv - R.output(gv) - return gv - - graph_model4 = fx.symbolic_trace(InterpolateNHWC4()) - with torch.no_grad(): - mod4 = from_fx(graph_model4, input_info_5d, default_image_layout="NDHWC") - tvm.ir.assert_structural_equal(mod4, expected_nhwc4) + verify_model( + Interpolate(), [(shape, "float32")], {}, builder.get(), default_image_layout=layout + ) def test_addmm(): @@ -6551,40 +6060,27 @@ def main( R.output(gv) return gv + @I.ir_module + class ExpectedNegative: + @R.function + def main(x: R.Tensor((3, 4), "float32")): + with R.dataflow(): + scaled = R.divide(x, R.const(10.0, "float32")) + rounded = R.round(scaled) + result = R.multiply(rounded, R.const(10.0, "float32")) + output = result + R.output(output) + return output + rounds = [ (0, Expected1), (2, Expected2), + (-1, ExpectedNegative), ] for decimals, expected in rounds: verify_model(Round(decimals), input_info, {}, expected) - # Test numerical accuracy with decimals - test_data = torch.tensor( - [ - [1.2345, 2.3456, 3.4567, 4.5678], - [5.6789, 6.7890, 7.8901, 8.9012], - [9.1234, 10.2345, 11.3456, 12.4567], - ] - ) - - for decimals in [0, 2]: - torch_model = Round(decimals) - graph_model = fx.symbolic_trace(torch_model) - with torch.no_grad(): - mod = from_fx(graph_model, input_info) - - target = tvm.target.Target("llvm") - ex = relax.build(mod, target) - vm = relax.VirtualMachine(ex, tvm.cpu()) - - torch_result = torch_model(test_data).numpy() - tvm_input = tvm.runtime.tensor(test_data.numpy()) - tvm_result = vm["main"](tvm_input).numpy() - - # Use relaxed tolerance due to floating-point precision in decimal operations - tvm.testing.assert_allclose(tvm_result, torch_result, rtol=1e-3, atol=1e-3) - if __name__ == "__main__": tvm.testing.main() From 15b2dc36b76c44e50e93cbb82a4e0c16bc1e6e23 Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Tue, 15 Sep 2026 20:58:09 -0400 Subject: [PATCH 3/3] Revert "[Tests][Frontend] Trim redundant PyTorch frontend coverage" This reverts commit f9036633d640acab2acd86e7626b488e46843c6d. --- .../test_frontend_from_exported_program.py | 1121 ++++++++++++++--- tests/python/relax/test_frontend_from_fx.py | 734 +++++++++-- 2 files changed, 1565 insertions(+), 290 deletions(-) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index cc58a02eba9e..c74573c49f14 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -62,38 +62,35 @@ def verify_model( tvm.ir.assert_structural_equal(mod, expected, map_free_vars=map_free_vars) -def verify_model_numerically( - torch_model, - example_args, - rtol=1e-7, - atol=1e-7, - *, - dynamic_shapes=None, - input_sets=None, - run_ep_decomposition=True, -): - """Build once and compare every output, including runs with different input shapes.""" - if not env.has_llvm(): - pytest.skip("need llvm") - exported_program = export(torch_model, args=example_args, dynamic_shapes=dynamic_shapes) - mod = from_exported_program(exported_program, run_ep_decomposition=run_ep_decomposition) - vm = relax.VirtualMachine(relax.build(mod, target="llvm"), tvm.cpu()) +def verify_model_numerically(torch_model, example_args, rtol=1e-7, atol=1e-7): + """Verify model by comparing numerical outputs between PyTorch and TVM.""" + with torch.no_grad(): + pytorch_output = torch_model(*example_args) - for args in (example_args,) if input_sets is None else input_sets: - tvm_args = [tvm.runtime.tensor(arg.detach().numpy()) for arg in args] - with torch.no_grad(): - expected = torch_model(*(arg.clone() for arg in args)) - expected = torch.utils._pytree.tree_leaves(expected) - actual = vm["main"](*tvm_args) - actual = [actual] if isinstance(actual, tvm.runtime.Tensor) else list(actual) - assert len(actual) == len(expected) - for actual_value, expected_value in zip(actual, expected): - actual_array = actual_value.numpy() - expected_array = expected_value.numpy() - assert actual_array.shape == expected_array.shape - assert actual_array.dtype == expected_array.dtype - np.testing.assert_allclose(actual_array, expected_array, rtol=rtol, atol=atol) - return mod + exported_program = export(torch_model, args=example_args) + mod = from_exported_program(exported_program) + target = tvm.target.Target("llvm") + ex = relax.build(mod, target) + vm = relax.VirtualMachine(ex, tvm.cpu()) + + tvm_args = [tvm.runtime.tensor(arg.numpy()) for arg in example_args] + tvm_output = vm["main"](*tvm_args) + + if hasattr(tvm_output, "numpy"): + tvm_output_np = tvm_output.numpy() + else: + tvm_output_np = tvm_output[0].numpy() + + pytorch_output_np = ( + pytorch_output.numpy() + if isinstance(pytorch_output, torch.Tensor) + else pytorch_output[0].numpy() + ) + + assert pytorch_output_np.shape == tvm_output_np.shape, ( + f"Shape mismatch: PyTorch {pytorch_output_np.shape} vs TVM {tvm_output_np.shape}" + ) + tvm.testing.assert_allclose(pytorch_output_np, tvm_output_np, rtol=rtol, atol=atol) operator_basic_unary = [ @@ -151,47 +148,96 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( verify_model(UnaryOp(), example_args, {}, expected) -@pytest.mark.parametrize( - "dtype, decimals", - [ - pytest.param(torch.float32, (0, 1, -1), id="ties-to-even"), - pytest.param(torch.float16, (4, 5, -5), id="float16-scaling"), - ], -) -def test_round_decimals(dtype, decimals): - class RoundDecimals(Module): - def forward(self, x): - return tuple(torch.round(x, decimals=d) for d in decimals) +def test_round_decimals(): + """torch.round(x, decimals) is exported as aten.round.decimals, which was missing + from the convert map (only round.default was registered) and made any explicit + decimals -- including decimals=0 -- fail with + "AssertionError: Unsupported function types ['round.decimals']". - # Negative decimals catch reciprocal-scaling errors; float16 catches overflow - # both in the scaled input (4) and in the scale itself (5 and -5). - x = torch.tensor([0.5, 1.5, 2.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25], dtype=dtype) - verify_model_numerically(RoundDecimals(), (x,), rtol=1e-6, atol=1e-6) + With the decimals overload registered, torch.round(x, decimals) must convert and + match PyTorch's round-half-to-even results, including negative decimals + (round(25, -1) == 20) where the scale-by-0.1 float precision path used to be wrong. + """ + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals -def test_round_decimals_large(): - """An overflowing scale must remain importable in either direction.""" + def forward(self, input): + return torch.round(input, decimals=self.decimals) - class RoundDecimals(Module): - def forward(self, x): - return torch.round(x, decimals=309), torch.round(x, decimals=-309) + # Half values exercise ties-to-even; 25/125/165 exercise the negative-decimals path. + x = torch.tensor( + [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25], dtype=torch.float32 + ) + for decimals in (0, 1, -1, -2): + verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,), rtol=1e-6, atol=1e-6) - @I.ir_module - class Expected: - @R.function - def main(x: R.Tensor((2,), "float32")): - with R.dataflow(): - scaled_up = R.multiply(x, R.const(float("inf"), "float32")) - rounded_up = R.round(scaled_up) - positive = R.divide(rounded_up, R.const(float("inf"), "float32")) - scaled_down = R.divide(x, R.const(float("inf"), "float32")) - rounded_down = R.round(scaled_down) - negative = R.multiply(rounded_down, R.const(float("inf"), "float32")) - result = (positive, negative) - R.output(result) - return result - verify_model(RoundDecimals(), (torch.ones(2),), {}, Expected) +def test_round_decimals_low_precision(): + """Scaling for low-precision inputs must happen in float32 and be cast back. + + 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 == 250000 + exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5 + already overflows float16 (the scale itself becomes inf), turning decimals=5 + and -5 into NaN. Upcasting the input to float32 keeps the scaling exact; the + rounded result is cast back to the input dtype. + """ + + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals + + def forward(self, input): + return torch.round(input, decimals=self.decimals) + + x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5], dtype=torch.float16) + # Positive decimals exercise the multiply-by-10**d overflow (4, 5); + # negative decimals exercise the 10**|d| scale overflowing float16 (-5). + for decimals in (2, 4, 5, -2, -4, -5): + verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,), rtol=1e-6, atol=1e-6) + + +def test_round_decimals_large(): + """A large |decimals| must import and run without OverflowError. + + The scale 10**|decimals| used to be built as an unbounded host Python int + before being handed to relax.const, whose int-to-float conversion raises + OverflowError ("int too large to convert to float") once |decimals| >= 309 + (10**309 already exceeds the float64 range). PyTorch accepts such decimals and + exports a valid aten.round.decimals node, so importing the exported program + must not crash on them. The scale is now built directly in the float dtype and + saturates to inf once it leaves the finite range, matching PyTorch, whose + all-NaN result here comes from the same inf scale. + """ + + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals + + def forward(self, input): + return torch.round(input, decimals=self.decimals) + + x = torch.tensor([0.5, 1.5, 25.0, -0.5, 0.0], dtype=torch.float32) + for decimals in (309, -309): + exported_program = export(RoundDecimalsModel(decimals).eval(), args=(x,)) + mod = from_exported_program(exported_program) # used to raise OverflowError here + ex = relax.build(mod, target="llvm") + vm = relax.VirtualMachine(ex, tvm.cpu()) + tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) + got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() + + # The scale overflows to inf, and IEEE arithmetic turns every element into + # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining + # (empty here) finite elements separately, since allclose fails on NaN. + expected = torch.round(x, decimals=decimals) + actual = torch.as_tensor(got) + assert torch.equal(torch.isnan(actual), torch.isnan(expected)) + finite = ~torch.isnan(expected) + assert torch.allclose(actual[finite], expected[finite], rtol=1e-6, atol=1e-6) operator_bool_unary = [ @@ -1006,18 +1052,10 @@ def main( verify_model(Atan2(), example_args, {}, expected) -@pytest.mark.parametrize( - "torch_op, relax_op", - [ - (torch.logical_and, R.logical_and), - (torch.logical_or, R.logical_or), - (torch.logical_xor, R.logical_xor), - ], -) -def test_logical_binary(torch_op, relax_op): - class LogicalBinary(Module): +def test_logical_and(): + class LogicalAnd(Module): def forward(self, lhs, rhs): - return torch_op(lhs, rhs) + return torch.logical_and(lhs, rhs) @tvm.script.ir_module class expected: @@ -1030,7 +1068,7 @@ def main( with R.dataflow(): lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") - lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(lv, lv1) + lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_and(lv, lv1) gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv2,) R.output(gv) return gv @@ -1039,7 +1077,7 @@ def main( torch.randn(1, 3, 10, 10, dtype=torch.float32), torch.randn(1, 3, 10, 10, dtype=torch.float32), ) - verify_model(LogicalBinary(), example_args, {}, expected) + verify_model(LogicalAnd(), example_args, {}, expected) def test_logical_not(): @@ -1065,6 +1103,62 @@ def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( verify_model(LogicalNot(), example_args, {}, expected) +def test_logical_or(): + class LogicalOr(Module): + def forward(self, lhs, rhs): + return torch.logical_or(lhs, rhs) + + @tvm.script.ir_module + class expected: + @R.function + def main( + lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + rhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + ) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")): + # block 0 + with R.dataflow(): + lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") + lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") + lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_or(lv, lv1) + gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv2,) + R.output(gv) + return gv + + example_args = ( + torch.randn(1, 3, 10, 10, dtype=torch.float32), + torch.randn(1, 3, 10, 10, dtype=torch.float32), + ) + verify_model(LogicalOr(), example_args, {}, expected) + + +def test_logical_xor(): + class LogicalXor(Module): + def forward(self, lhs, rhs): + return torch.logical_xor(lhs, rhs) + + @tvm.script.ir_module + class expected: + @R.function + def main( + lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + rhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + ) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")): + # block 0 + with R.dataflow(): + lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") + lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") + lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_xor(lv, lv1) + gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv2,) + R.output(gv) + return gv + + example_args = ( + torch.randn(1, 3, 10, 10, dtype=torch.float32), + torch.randn(1, 3, 10, 10, dtype=torch.float32), + ) + verify_model(LogicalXor(), example_args, {}, expected) + + def test_pow_integer(): class Pow(Module): def forward(self, input): @@ -2508,6 +2602,15 @@ def __init__(self): def forward(self, input): return self.conv(input) + class ConvTranspose1d1Func(Module): + def __init__(self): + super().__init__() + self.weight = torch.randn(size=[6, 6, 3]) + self.bias = torch.randn(size=[6]) + + def forward(self, input): + return torch.nn.functional.conv_transpose1d(input, self.weight, self.bias) + @tvm.script.ir_module class expected1: @R.function @@ -2575,6 +2678,10 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) + model = ConvTranspose1d1Func() + binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} + verify_model(model, example_args, binding, expected1) + model = ConvTranspose1d2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -2589,6 +2696,15 @@ def __init__(self): def forward(self, input): return self.conv(input) + class ConvTranspose2d1Func(Module): + def __init__(self): + super().__init__() + self.weight = torch.randn(size=[3, 3, 7, 7]) + self.bias = torch.randn(size=[3]) + + def forward(self, input): + return torch.nn.functional.conv_transpose2d(input, self.weight, self.bias) + @tvm.script.ir_module class expected1: @R.function @@ -2656,6 +2772,10 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) + model = ConvTranspose2d1Func() + binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} + verify_model(model, example_args, binding, expected1) + model = ConvTranspose2d2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -2670,6 +2790,15 @@ def __init__(self): def forward(self, input): return self.conv(input) + class Conv1D1Func(Module): + def __init__(self): + super().__init__() + self.weight = torch.randn(size=[6, 3, 7]) + self.bias = torch.randn(size=[6]) + + def forward(self, input): + return torch.nn.functional.conv1d(input, self.weight, self.bias) + @tvm.script.ir_module class expected1: @R.function @@ -2735,6 +2864,10 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) + model = Conv1D1Func() + binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} + verify_model(model, example_args, binding, expected1) + model = Conv1D2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -2749,6 +2882,15 @@ def __init__(self): def forward(self, input): return self.conv(input) + class Conv2D1Func(Module): + def __init__(self): + super().__init__() + self.weight = torch.randn(size=[6, 3, 7, 7]) + self.bias = torch.randn(size=[6]) + + def forward(self, input): + return torch.nn.functional.conv2d(input, self.weight, self.bias) + @tvm.script.ir_module class expected1: @R.function @@ -2814,6 +2956,10 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) + model = Conv2D1Func() + binding = {"w1": model.weight.numpy(), "w2": model.bias.numpy()} + verify_model(model, example_args, binding, expected1) + model = Conv2D2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -2828,6 +2974,15 @@ def __init__(self): def forward(self, input): return self.conv(input) + class Conv3D1Func(Module): + def __init__(self): + super().__init__() + self.weight = torch.randn(size=[6, 3, 7, 7, 7]) + self.bias = torch.randn(size=[6]) + + def forward(self, input): + return torch.nn.functional.conv3d(input, self.weight, self.bias) + @tvm.script.ir_module class expected1: @R.function @@ -2893,6 +3048,10 @@ def main( binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) + model = Conv3D1Func() + binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} + verify_model(model, example_args, binding, expected1) + model = Conv3D2() binding = {"w1": model.conv.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -3261,7 +3420,19 @@ def main( def test_einsum_repeated_subscript(): - """Decomposed diagonal extraction must lower to a single einsum.""" + """einsum with repeated subscripts (diagonal / trace) on the default + decomposition path. + + ``run_decompositions`` (default) lowers repeated-subscript einsum to + ``aten.diagonal`` + ``permute`` (+ ``sum`` for the trace), which the + frontend converts with the ``_diagonal`` lowering. For the zero-offset + square case (e.g. ``torch.einsum("ii->i")`` on an ``N x N`` input) the + frontend emits a single repeated-subscript einsum that reads the diagonal + directly; otherwise it permutes the diagonal dims to the trailing two axes, + slices each to the diagonal length, and runs an einsum ``...zz->...z``. + This used to raise ``AssertionError: Unsupported function types + ['diagonal.default']``. + """ class EinsumDiag(Module): def __init__(self): @@ -3285,20 +3456,118 @@ def main(x: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((3,), dtype=" example_args = (torch.randn(3, 3, dtype=torch.float32),) verify_model(EinsumDiag(), example_args, {}, Expected) - class BatchedDiagonal(Module): + class TraceEinsum(Module): + def forward(self, x): + return torch.einsum("ii->", x) + + class BatchedDiagEinsum(Module): + def forward(self, x): + return torch.einsum("...ii->...i", x) + + class AttentionEinsum(Module): + def forward(self, x, y): + return torch.einsum("abca,abcb->c", x, y) + + verify_model_numerically(TraceEinsum(), (torch.randn(4, 4),)) + verify_model_numerically(BatchedDiagEinsum(), (torch.randn(2, 3, 3),)) + verify_model_numerically(AttentionEinsum(), (torch.randn(3, 3, 4, 3), torch.randn(3, 3, 4, 3))) + + class DirectDiagonal(Module): + def __init__(self): + super().__init__() + self.offset = 1 + def forward(self, x): - return torch.einsum("...ii->...i", x), torch.einsum("...ii->...", x) + return torch.diagonal(x, self.offset, 0, 1) - verify_model_numerically(BatchedDiagonal(), (torch.arange(18.0).reshape(2, 3, 3),)) + class DirectTrace(Module): + def forward(self, x): + return torch.trace(x) + + verify_model_numerically(DirectDiagonal(), (torch.randn(3, 4),)) + verify_model_numerically(DirectTrace(), (torch.randn(4, 4),)) + # Out-of-range offsets (|offset| >= max(extent1, extent2)) are valid in + # PyTorch and yield an empty diagonal of shape (0,); the lowering must + # clamp the diagonal length to zero instead of producing negative slice + # extents or a wrong non-empty shape. + class DirectDiagonalOutOfRange(Module): + def __init__(self, offset): + super().__init__() + self.offset = offset -def test_diagonal_offsets(): - class Diagonal(Module): def forward(self, x): - # Non-square input, both offset signs, and one empty result per sign. - return tuple(torch.diagonal(x, offset, 0, 1) for offset in (1, -1, 4, -3)) + return torch.diagonal(x, self.offset, 0, 1) - verify_model_numerically(Diagonal(), (torch.arange(12.0).reshape(3, 4),)) + for offset in [4, 5, 6, -3, -4, -5, -6]: + verify_model_numerically(DirectDiagonalOutOfRange(offset), (torch.randn(3, 4),)) + + +def test_einsum_diagonal_lowers_without_full_size_intermediate(): + """Regression test: a zero-offset square diagonal must not materialize + full-size intermediates. + + ``torch.einsum("ii->i")`` on an ``N x N`` input is decomposed to + ``aten.diagonal`` by ``run_decompositions``. Lowering that diagonal by + permuting the diagonal dims to the trailing axes, slicing each to the + diagonal length, and running the ``...zz->...z`` einsum materializes three + full-size ``N x N`` intermediates (an identity permute and two identity + strided slices) and hence three O(N^2) copy loops before the final O(N) + diagonal loop. The ``_diagonal`` fast path instead emits a single + repeated-subscript einsum that reads the diagonal directly, so no full-size + intermediate exists in the frontend graph (and therefore neither in the + lowered TIR). Assert that every intermediate produced by a call is at most + O(N), both before and after legalization. + """ + + class EinsumDiag(Module): + def forward(self, x): + return torch.einsum("ii->i", x) + + n = 8 + exported_program = export(EinsumDiag(), args=(torch.randn(n, n),)) + mod = from_exported_program(exported_program) + + def rank2_call_results(ir_mod): + """Names of calls whose result is a rank-2 (full-size) tensor.""" + results = [] + for func in ir_mod.functions.values(): + if not isinstance(func, relax.Function): + continue + for block in func.body.blocks: + for binding in block.bindings: + if not ( + isinstance(binding.value, relax.Call) + and isinstance(binding.value.op, tvm.ir.Op) + ): + continue + if isinstance(binding.var.ty, relax.TensorType) and binding.var.ty.ndim == 2: + results.append(binding.value.op.name) + return results + + # The diagonal must be the only full-size (N x N) tensor touched: it is the + # function input read directly by a single repeated-subscript einsum. No + # call may produce a rank-2 intermediate. + assert rank2_call_results(mod) == [] + + # Sanity check that the graph really performs the diagonal: exactly one + # einsum on the N x N input producing an N-vector. + einsum_calls = [] + for block in mod["main"].body.blocks: + for binding in block.bindings: + if ( + isinstance(binding.value, relax.Call) + and isinstance(binding.value.op, tvm.ir.Op) + and binding.value.op.name == "relax.einsum" + ): + einsum_calls.append(binding.var) + assert len(einsum_calls) == 1 + assert einsum_calls[0].ty.ndim == 1 + + # Legalize and check again on the lowered graph. + with tvm.target.Target("llvm"): + lowered = relax.transform.LegalizeOps()(mod) + assert rank2_call_results(lowered) == [] def test_outer(): @@ -3492,6 +3761,15 @@ def __init__(self): def forward(self, input): return self.linear(input) + class Dense1Func(Module): + def __init__(self): + super().__init__() + self.weight = torch.randn(size=[7, 10]) + self.bias = torch.randn(size=[7]) + + def forward(self, input): + return torch.nn.functional.linear(input, self.weight, self.bias) + @tvm.script.ir_module class expected1: @R.function @@ -3546,6 +3824,10 @@ def main( binding = {"w1": model.linear.weight.detach().numpy(), "w2": model.linear.bias.detach().numpy()} verify_model(model, example_args, binding, expected1) + model = Dense1Func() + binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()} + verify_model(model, example_args, binding, expected1) + model = Dense2() binding = {"w1": model.linear.weight.detach().numpy()} verify_model(model, example_args, binding, expected2) @@ -3641,6 +3923,13 @@ def __init__(self): def forward(self, input): return self.pool(input) + class MaxPool2d_functional(Module): + def __init__(self): + super().__init__() + + def forward(self, input): + return torch.nn.functional.max_pool2d(input, kernel_size=[1, 1]) + @tvm.script.ir_module class expected1: @R.function @@ -3738,6 +4027,7 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),) verify_model(MaxPool2d(), example_args, {}, expected1) + verify_model(MaxPool2d_functional(), example_args, {}, expected1) verify_model(MaxPool2d2(), example_args, {}, expected2) verify_model(MaxPool2d3(), example_args, {}, expected3) @@ -3751,6 +4041,13 @@ def __init__(self): def forward(self, input): return self.pool(input) + class MaxPool3d_functional(Module): + def __init__(self): + super().__init__() + + def forward(self, input): + return torch.nn.functional.max_pool3d(input, kernel_size=[1, 1, 1]) + @tvm.script.ir_module class expected1: @R.function @@ -3852,6 +4149,7 @@ def main(input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32")) -> R.Tuple( # Verify the models with expected IR modules verify_model(MaxPool3d(), example_args1, {}, expected1) + verify_model(MaxPool3d_functional(), example_args1, {}, expected1) verify_model(MaxPool3d2(), example_args2, {}, expected2) verify_model(MaxPool3d3(), example_args3, {}, expected3) @@ -5518,25 +5816,37 @@ def forward(self, x): example_args = (torch.randn(1, 4, 3, dtype=torch.float32),) tokens = torch.export.Dim("tokens", min=1, max=8) - mod = verify_model_numerically( + exported_program = export( DynamicShapeOps(), - example_args, + args=example_args, dynamic_shapes={"x": {1: tokens}}, - input_sets=[(torch.randn(1, token_count, 3),) for token_count in (4, 6)], - rtol=0, - atol=0, ) + mod = from_exported_program(exported_program) + script = mod.script() assert "R.tensor_to_shape" in script assert "R.shape_to_tensor" in script + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + for token_count in (4, 6): + torch_input = torch.randn(1, token_count, 3) + expected = DynamicShapeOps()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @pytest.mark.parametrize( ("item_dtype", "item_shape"), [ + (torch.int8, ()), + (torch.uint8, ()), + (torch.int16, ()), (torch.int32, ()), - (torch.uint8, (1,)), + (torch.int64, ()), + (torch.int64, (1,)), (torch.int64, (1, 1)), ], ) @@ -5557,15 +5867,21 @@ def forward(self, x): rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) example_args = (torch.randn(3, 4, dtype=torch.float32),) - verify_model_numerically( + exported_program = export( DynamicItem(), - example_args, - rtol=0, - atol=0, + args=example_args, dynamic_shapes={"x": {0: rows, 1: columns}}, - input_sets=[(torch.randn(shape),) for shape in ((3, 4), (5, 2))], - run_ep_decomposition=False, ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 4), (5, 2)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicItem()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -5584,9 +5900,15 @@ def forward(self, x, signed, unsigned, wide): torch.tensor([200], dtype=torch.uint8), torch.tensor(1 << 40, dtype=torch.int64), ) - verify_model_numerically( - RuntimeItems(), example_args, rtol=0, atol=0, run_ep_decomposition=False - ) + exported_program = export(RuntimeItems(), args=example_args) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + expected = RuntimeItems()(*example_args) + actual = vm["main"](*(tvm.runtime.tensor(arg.numpy()) for arg in example_args)) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -5603,15 +5925,26 @@ def forward(self, x): example_args = (torch.randn(3, 4, dtype=torch.float32),) rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) - verify_model_numerically( + exported_program = export( DynamicFills(), - example_args, - rtol=0, - atol=0, + args=example_args, dynamic_shapes={"x": {0: rows, 1: columns}}, - input_sets=[(torch.randn(shape),) for shape in ((3, 4), (5, 2))], - run_ep_decomposition=False, ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 4), (5, 2)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicFills()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + actual_arrays = [value.numpy() for value in actual] + + assert actual_arrays[0].dtype == np.dtype("int64") + assert actual_arrays[1].dtype == np.dtype("float64") + for actual_value, expected_value in zip(actual_arrays, expected): + assert actual_value.dtype == expected_value.numpy().dtype + np.testing.assert_array_equal(actual_value, expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -5635,15 +5968,21 @@ def forward(self, x): rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) example_args = (torch.randn(3, 4, dtype=torch.float32),) - verify_model_numerically( + exported_program = export( DynamicBooleanFills(), - example_args, - rtol=0, - atol=0, + args=example_args, dynamic_shapes={"x": {0: rows, 1: columns}}, - input_sets=[(torch.randn(shape),) for shape in ((3, 3), (3, 4))], - run_ep_decomposition=False, ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 3), (3, 4)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicBooleanFills()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -5666,15 +6005,21 @@ def forward(self, x): rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) example_args = (torch.randn(3, 4, dtype=torch.float32),) - verify_model_numerically( + exported_program = export( DynamicScalarArithmetic(), - example_args, - rtol=0, - atol=0, + args=example_args, dynamic_shapes={"x": {0: rows, 1: columns}}, - input_sets=[(torch.randn(shape),) for shape in ((3, 4), (5, 2))], - run_ep_decomposition=False, ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 4), (5, 2)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicScalarArithmetic()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -5694,7 +6039,16 @@ def forward(self, x): ) example_args = (torch.randn(2, 3, dtype=torch.float32),) - verify_model_numerically(FullLike(), example_args, rtol=0, atol=0, run_ep_decomposition=False) + exported_program = export(FullLike(), args=example_args) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + expected = FullLike()(*example_args) + actual = vm["main"](tvm.runtime.tensor(example_args[0].numpy())) + for actual_value, expected_value in zip(actual, expected): + assert actual_value.numpy().dtype == expected_value.numpy().dtype + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) def test_split(): @@ -5732,7 +6086,16 @@ def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( def test_split_int_split_size(): - """A non-divisible split size denotes chunk length, not section count.""" + """x.split(int, dim) must produce chunks of size `split_size` (the last one + smaller when the dimension is not divisible), matching PyTorch. + + The frontend used to convert the int per-chunk size into a section count and + pass it as relax.op.split's int argument, which means "split into N equal + sections"; that yields wrong chunk shapes whenever + ceil(D / ceil(D / split_size)) != split_size (e.g. split_size > D/2). The + int branch now builds cumulative cut positions, the same as the list/tuple + form. + """ class Split6(Module): def forward(self, input): @@ -5763,6 +6126,43 @@ def main(input: R.Tensor((10,), dtype="float32")) -> R.Tuple( verify_model(Split6(), example_args, {}, Expected) # Differential check against native PyTorch for non-divisible sizes and dims. + class SplitModel(Module): + def __init__(self, split_size, dim): + super().__init__() + self.split_size = split_size + self.dim = dim + + def forward(self, input): + return input.split(self.split_size, dim=self.dim) + + def run_tvm(model, args): + exported_program = export(model, args=args) + mod = from_exported_program(exported_program) + ex = relax.build(mod, target="llvm") + vm = relax.VirtualMachine(ex, tvm.cpu()) + out = vm["main"](*[tvm.runtime.tensor(a.numpy()) for a in args]) + if hasattr(out, "numpy"): + return [out.numpy()] + return [o.numpy() for o in out] + + for shape, split_size, dim in [ + ((10,), 6, 0), + ((10,), 7, 0), + ((10,), 8, 0), + ((10,), 9, 0), + ((12,), 7, 0), + ((12, 8), 5, 1), + ((3, 10), 6, -1), + ]: + x = torch.arange(1, int(np.prod(shape)) + 1, dtype=torch.float32).reshape(shape) + refs = [r.numpy() for r in x.split(split_size, dim)] + outs = run_tvm(SplitModel(split_size, dim), (x,)) + assert [r.shape for r in refs] == [o.shape for o in outs], ( + f"split shape={shape} s={split_size} dim={dim}: " + f"torch {[r.shape for r in refs]} vs tvm {[o.shape for o in outs]}" + ) + for r, o in zip(refs, outs): + tvm.testing.assert_allclose(o, r, rtol=1e-7, atol=1e-7) def test_squeeze(): @@ -6267,15 +6667,21 @@ def forward(self, x): rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) example_args = (torch.randn(3, 4, dtype=torch.float32),) - verify_model_numerically( + exported_program = export( DynamicMaskedFills(), - example_args, - rtol=0, - atol=0, + args=example_args, dynamic_shapes={"x": {0: rows, 1: columns}}, - input_sets=[(torch.randn(shape),) for shape in ((3, 4), (5, 2))], - run_ep_decomposition=False, ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 4), (5, 2)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicMaskedFills()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") @@ -6301,9 +6707,16 @@ def forward(self, x, mask): torch.arange(6, dtype=dtype).reshape(2, 3), torch.tensor([[True, False, True], [False, True, False]]), ) - verify_model_numerically( - MaskedFills(), example_args, rtol=0, atol=0, run_ep_decomposition=False - ) + exported_program = export(MaskedFills(), args=example_args) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + expected = MaskedFills()(*example_args) + actual = vm["main"](*(tvm.runtime.tensor(arg.numpy()) for arg in example_args)) + for actual_value, expected_value in zip(actual, expected): + assert actual_value.numpy().dtype == expected_value.numpy().dtype + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) def test_masked_select(): @@ -6857,10 +7270,18 @@ class Gather0(Module): def forward(self, data, indices): return torch.gather(data, 0, indices) + class Gather1(Module): + def forward(self, data, indices): + return torch.gather(data, 1, indices) + class Gather2(Module): def forward(self, data, indices): return torch.gather(data, -1, indices) + class Gather3(Module): + def forward(self, data, indices): + return torch.gather(data, -2, indices) + @tvm.script.ir_module class Expected0: @R.function @@ -6874,6 +7295,19 @@ def main( R.output(gv) return gv + @tvm.script.ir_module + class Expected1: + @R.function + def main( + inp_0: R.Tensor((2, 3), dtype="float32"), + inp_1: R.Tensor((2, 3), dtype="int64"), + ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=1) + gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,) + R.output(gv) + return gv + @tvm.script.ir_module class Expected2: @R.function @@ -6887,13 +7321,28 @@ def main( R.output(gv) return gv + @tvm.script.ir_module + class Expected3: + @R.function + def main( + inp_0: R.Tensor((2, 3), dtype="float32"), + inp_1: R.Tensor((2, 3), dtype="int64"), + ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=-2) + gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,) + R.output(gv) + return gv + example_args = ( torch.randn(2, 3, dtype=torch.float32), torch.randint(0, 3, (2, 3), dtype=torch.int64), ) verify_model(Gather0(), example_args, {}, Expected0) + verify_model(Gather1(), example_args, {}, Expected1) verify_model(Gather2(), example_args, {}, Expected2) + verify_model(Gather3(), example_args, {}, Expected3) def test_index_put(): @@ -6956,10 +7405,112 @@ def main( return gv # Test case 3: 3D input + class IndexPut3D(Module): + def forward(self, data, indices_0, indices_1, indices_2, values): + indices_tuple = (indices_0, indices_1, indices_2) + return data.index_put_(indices_tuple, values, accumulate=False) + + example_args_3d = ( + torch.randn(16, 32, 64, dtype=torch.float32), + torch.randint(0, 16, (128,), dtype=torch.int64), + torch.randint(0, 32, (128,), dtype=torch.int64), + torch.randint(0, 64, (128,), dtype=torch.int64), + torch.randn(128, dtype=torch.float32), + ) + + @I.ir_module + class Expected3D: + @R.function + def main( + data: R.Tensor((16, 32, 64), dtype="float32"), + indices_0: R.Tensor((128,), dtype="int64"), + indices_1: R.Tensor((128,), dtype="int64"), + indices_2: R.Tensor((128,), dtype="int64"), + values: R.Tensor((128,), dtype="float32"), + ) -> R.Tuple(R.Tensor((16, 32, 64), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((16, 32, 64), dtype="float32") = R.index_put( + data, (indices_0, indices_1, indices_2), values, accumulate=False + ) + gv: R.Tuple(R.Tensor((16, 32, 64), dtype="float32")) = (lv,) + R.output(gv) + return gv # Test case 4: 4D input + class IndexPut4D(Module): + def forward(self, data, indices_0, indices_1, indices_2, indices_3, values): + indices_tuple = (indices_0, indices_1, indices_2, indices_3) + return data.index_put_(indices_tuple, values, accumulate=False) + + example_args_4d = ( + torch.randn(8, 16, 32, 64, dtype=torch.float32), + torch.randint(0, 8, (128,), dtype=torch.int64), + torch.randint(0, 16, (128,), dtype=torch.int64), + torch.randint(0, 32, (128,), dtype=torch.int64), + torch.randint(0, 64, (128,), dtype=torch.int64), + torch.randn(128, dtype=torch.float32), + ) + + @I.ir_module + class Expected4D: + @R.function + def main( + data: R.Tensor((8, 16, 32, 64), dtype="float32"), + indices_0: R.Tensor((128,), dtype="int64"), + indices_1: R.Tensor((128,), dtype="int64"), + indices_2: R.Tensor((128,), dtype="int64"), + indices_3: R.Tensor((128,), dtype="int64"), + values: R.Tensor((128,), dtype="float32"), + ) -> R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((8, 16, 32, 64), dtype="float32") = R.index_put( + data, + (indices_0, indices_1, indices_2, indices_3), + values, + accumulate=False, + ) + gv: R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")) = (lv,) + R.output(gv) + return gv # Test case 5: 5D input + class IndexPut5D(Module): + def forward(self, data, indices_0, indices_1, indices_2, indices_3, indices_4, values): + indices_tuple = (indices_0, indices_1, indices_2, indices_3, indices_4) + return data.index_put_(indices_tuple, values, accumulate=False) + + example_args_5d = ( + torch.randn(4, 8, 16, 32, 64, dtype=torch.float32), + torch.randint(0, 4, (128,), dtype=torch.int64), + torch.randint(0, 8, (128,), dtype=torch.int64), + torch.randint(0, 16, (128,), dtype=torch.int64), + torch.randint(0, 32, (128,), dtype=torch.int64), + torch.randint(0, 64, (128,), dtype=torch.int64), + torch.randn(128, dtype=torch.float32), + ) + + @I.ir_module + class Expected5D: + @R.function + def main( + data: R.Tensor((4, 8, 16, 32, 64), dtype="float32"), + indices_0: R.Tensor((128,), dtype="int64"), + indices_1: R.Tensor((128,), dtype="int64"), + indices_2: R.Tensor((128,), dtype="int64"), + indices_3: R.Tensor((128,), dtype="int64"), + indices_4: R.Tensor((128,), dtype="int64"), + values: R.Tensor((128,), dtype="float32"), + ) -> R.Tuple(R.Tensor((4, 8, 16, 32, 64), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((4, 8, 16, 32, 64), dtype="float32") = R.index_put( + data, + (indices_0, indices_1, indices_2, indices_3, indices_4), + values, + accumulate=False, + ) + gv: R.Tuple(R.Tensor((4, 8, 16, 32, 64), dtype="float32")) = (lv,) + R.output(gv) + return gv # Test case 6: 2D input with multi-dimensional index (broadcasting) # This tests the multi-dimensional index support with broadcasting @@ -7109,6 +7660,9 @@ def main(x: R.Tensor((2, 10), dtype="float32")) -> R.Tuple( # Run verification for each case verify_model(IndexPut1D(), example_args_1d, {}, Expected1D) verify_model(IndexPut2D(), example_args_2d, {}, Expected2D) + verify_model(IndexPut3D(), example_args_3d, {}, Expected3D) + verify_model(IndexPut4D(), example_args_4d, {}, Expected4D) + verify_model(IndexPut5D(), example_args_5d, {}, Expected5D) verify_model(IndexPutBroadcast1D(), example_args_broadcast1, {}, ExpectedBroadcast1D) verify_model(IndexPutBroadcast2D(), example_args_broadcast2, {}, ExpectedBroadcast2D) verify_model(IndexPutBroadcast3D(), example_args_broadcast3d, {}, ExpectedBroadcast3D) @@ -7142,6 +7696,40 @@ def forward(self, x, buf, idx): ) +def test_m4d_diag_index_put_tuple_output_regression(): + class M4D(Module): + def forward(self, x): + b, k, n = 2, 3, 5 + buf = x.new_zeros(b, k, n, n) + idx = torch.arange(n, device=x.device) + + diag = buf[..., idx, idx] + diag = torch.nn.functional.elu(diag) + 1.0 + 1e-8 + buf[..., idx, idx] = diag + + return x[..., :1], buf + + ex_in = torch.zeros(2, 3, 5, dtype=torch.float32) + exported_program = export(M4D().eval(), args=(ex_in,)) + + exported_targets = [str(getattr(n, "target", "")) for n in exported_program.graph.nodes] + assert any("index_put" in target for target in exported_targets) + + # Regression focus: importing this graph should not segfault at Tuple construction. + mod = from_exported_program(exported_program) + ret_ty = mod["main"].ret_ty + assert isinstance(ret_ty, relax.TupleType) + + tensor_fields = [f for f in ret_ty.fields if isinstance(f, relax.TensorType)] + assert len(tensor_fields) >= 2 + # x: (2, 3, 5) → x[..., :1]: (2, 3, 1) + assert any(len(f.shape) == 3 and int(f.shape[-1]) == 1 for f in tensor_fields) + # buf: (2, 3, 5, 5) → 4-D with spatial dims 5x5 + assert any( + len(f.shape) == 4 and int(f.shape[-2]) == 5 and int(f.shape[-1]) == 5 for f in tensor_fields + ) + + def test_index_put_mutation_through_alias_regression(): class IndexPutAlias(Module): def forward(self, x, idx, values): @@ -7190,6 +7778,10 @@ class Flip0(Module): def forward(self, data): return torch.flip(data, [0]) + class Flip1(Module): + def forward(self, data): + return torch.flip(data, [1]) + @tvm.script.ir_module class Expected0: @R.function @@ -7202,9 +7794,22 @@ def main( R.output(gv) return gv + @tvm.script.ir_module + class Expected1: + @R.function + def main( + inp_0: R.Tensor((2, 2), dtype="float32"), + ) -> R.Tuple(R.Tensor((2, 2), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((2, 2), dtype="float32") = R.flip(inp_0, axis=1) + gv: R.Tuple(R.Tensor((2, 2), dtype="float32")) = (lv,) + R.output(gv) + return gv + example_args = (torch.randn(2, 2, dtype=torch.float32),) verify_model(Flip0(), example_args, {}, Expected0) + verify_model(Flip1(), example_args, {}, Expected1) def test_flip_multi_axis(): @@ -7481,7 +8086,8 @@ def main( verify_model(Bucketize(), (input_tensor, boundaries), {}, Expected) -@pytest.mark.parametrize("right, out_int32", [(False, False), (True, True)]) +@pytest.mark.parametrize("right", [False, True]) +@pytest.mark.parametrize("out_int32", [False, True]) def test_bucketize_numerically(right, out_int32): class Bucketize(Module): def forward(self, input_tensor, boundaries): @@ -7999,34 +8605,42 @@ def main( verify_model(SparseMatrixMultiply(), example_args, {}, Expected) -@pytest.mark.parametrize("rnn_type", [nn.LSTM, nn.GRU, nn.RNN], ids=["lstm", "gru", "rnn-tanh"]) -@pytest.mark.parametrize( - "batch_first, bidirectional", - [(True, False), (False, True)], - ids=["batch-first", "bidirectional"], -) -def test_recurrent(rnn_type, batch_first, bidirectional): - class Recurrent(Module): - def __init__(self): +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_lstm(): + class LSTM(nn.Module): + def __init__(self, input_size, hidden_size, batch_first, bidirectional): super().__init__() - self.rnn = rnn_type(3, 4, batch_first=batch_first, bidirectional=bidirectional) + self.lstm = nn.LSTM( + input_size=input_size, + hidden_size=hidden_size, + num_layers=1, + batch_first=batch_first, + bidirectional=bidirectional, + ) def forward(self, x): - output, state = self.rnn(x) - return (output, state) if rnn_type is nn.RNN else output - - # Exercise both layouts and direction counts without repeating their product. - # Retain the RNN hidden-state check alongside its sequence output. - with torch.random.fork_rng(devices=[]): - torch.manual_seed(42) - x = torch.randn(2, 3, 3) if batch_first else torch.randn(3, 2, 3) - verify_model_numerically( - Recurrent(), - (x,), - rtol=1e-4, - atol=1e-5, - run_ep_decomposition=rnn_type is not nn.RNN, - ) + y, _ = self.lstm(x) + return y + + # Unidirectional LSTM with batch_first=True + torch.manual_seed(42) + x = torch.randn(2, 3, 4, dtype=torch.float32) + verify_model_numerically(LSTM(4, 8, batch_first=True, bidirectional=False), (x,)) + + # Unidirectional LSTM with batch_first=False + torch.manual_seed(43) + x2 = torch.randn(4, 2, 3, dtype=torch.float32) + verify_model_numerically(LSTM(3, 6, batch_first=False, bidirectional=False), (x2,)) + + # Bidirectional LSTM with batch_first=True + torch.manual_seed(44) + x3 = torch.randn(2, 3, 4, dtype=torch.float32) + verify_model_numerically(LSTM(4, 8, batch_first=True, bidirectional=True), (x3,)) + + # Bidirectional LSTM with batch_first=False + torch.manual_seed(45) + x4 = torch.randn(4, 2, 3, dtype=torch.float32) + verify_model_numerically(LSTM(3, 6, batch_first=False, bidirectional=True), (x4,)) def test_tensor_none_tuple(): @@ -8051,6 +8665,96 @@ def main(x: R.Tensor((3,), dtype="float32")) -> R.Tuple( verify_model(TensorNoneModel(), example_args, {}, Expected) +def test_gru(): + class GRU(nn.Module): + def __init__(self, input_size, hidden_size, batch_first, bidirectional): + super().__init__() + self.gru = nn.GRU( + input_size=input_size, + hidden_size=hidden_size, + num_layers=1, + batch_first=batch_first, + bidirectional=bidirectional, + ) + + def forward(self, x): + y, _ = self.gru(x) + return y + + cases = [ + (42, (2, 3, 4), 4, 8, True, False), + (43, (4, 2, 3), 3, 6, False, False), + (44, (2, 3, 4), 4, 5, True, True), + (45, (4, 2, 3), 3, 4, False, True), + ] + for seed, shape, input_size, hidden_size, batch_first, bidirectional in cases: + torch.manual_seed(seed) + x = torch.randn(*shape, dtype=torch.float32) + verify_model_numerically( + GRU(input_size, hidden_size, batch_first, bidirectional), + (x,), + rtol=1e-4, + atol=1e-5, + ) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_rnn_tanh(): + target = tvm.target.Target("llvm") + + def _check(rnn_kwargs, x_shape, seed): + class RNNWithState(nn.Module): + def __init__(self): + super().__init__() + self.rnn = nn.RNN(nonlinearity="tanh", num_layers=1, **rnn_kwargs) + + def forward(self, x): + output, h_n = self.rnn(x) + return output, h_n + + torch.manual_seed(seed) + x = torch.randn(*x_shape, dtype=torch.float32) + model = RNNWithState() + with torch.no_grad(): + pt_out, pt_hn = model(x) + + exported_program = export(model, args=(x,)) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + ex = relax.build(mod, target) + vm = relax.VirtualMachine(ex, tvm.cpu()) + tvm_outputs = vm["main"](tvm.runtime.tensor(x.numpy())) + tvm_out_np = tvm_outputs[0].numpy() + tvm_hn_np = tvm_outputs[1].numpy() + + assert pt_out.shape == tvm_out_np.shape, ( + f"output shape mismatch: PyTorch {tuple(pt_out.shape)} vs TVM {tvm_out_np.shape}" + ) + assert pt_hn.shape == tvm_hn_np.shape, ( + f"h_n shape mismatch: PyTorch {tuple(pt_hn.shape)} vs TVM {tvm_hn_np.shape}" + ) + tvm.testing.assert_allclose(pt_out.numpy(), tvm_out_np, rtol=1e-4, atol=1e-5) + tvm.testing.assert_allclose(pt_hn.numpy(), tvm_hn_np, rtol=1e-4, atol=1e-5) + + # batch_first, unidirectional + _check( + {"input_size": 4, "hidden_size": 8, "batch_first": True, "bidirectional": False}, + (2, 3, 4), + seed=42, + ) + # seq-first (batch_first=False), unidirectional + _check( + {"input_size": 3, "hidden_size": 6, "batch_first": False, "bidirectional": False}, + (4, 2, 3), + seed=43, + ) + # bidirectional, batch_first + _check( + {"input_size": 4, "hidden_size": 8, "batch_first": True, "bidirectional": True}, + (2, 3, 4), + seed=44, + ) + + def test_dynamic_shape_with_range_constraints(): class DynamicModel(torch.nn.Module): def forward(self, x1, x2): @@ -8508,6 +9212,10 @@ class UpsampleNearest2dScale(Module): def forward(self, input): return torch.nn.functional.interpolate(input, scale_factor=2.0, mode="nearest") + class UpsampleNearest2dSize(Module): + def forward(self, input): + return torch.nn.functional.interpolate(input, size=(20, 20), mode="nearest") + example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),) @tvm.script.ir_module @@ -8528,7 +9236,26 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( R.output(gv) return gv + @tvm.script.ir_module + class expected_size: + @R.function + def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple( + R.Tensor((1, 3, 20, 20), dtype="float32") + ): + with R.dataflow(): + lv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.image.resize2d( + input_1, + size=(20, 20), + layout="NCHW", + method="nearest_neighbor", + coordinate_transformation_mode="half_pixel", + ) + gv: R.Tuple(R.Tensor((1, 3, 20, 20), dtype="float32")) = (lv,) + R.output(gv) + return gv + verify_model(UpsampleNearest2dScale(), example_args, {}, expected_scale) + verify_model(UpsampleNearest2dSize(), example_args, {}, expected_size) def test_from_exported_program_sparse_csr_buffer(): @@ -8665,9 +9392,9 @@ def main( ) -@pytest.mark.parametrize("compare", [operator.eq, operator.ne], ids=["equal", "not-equal"]) -def test_cond_shape_comparison(compare): - class CondShapeModel(Module): +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_cond_shape_equality_predicate(): + class CondShapeEqualityModel(Module): def forward(self, x): def true_fn(x): return x + 1.0 @@ -8675,18 +9402,52 @@ def true_fn(x): def false_fn(x): return x - 1.0 - return torch.cond(compare(x.shape[0], x.shape[1]), true_fn, false_fn, (x,)) + return torch.cond(x.shape[0] == x.shape[1], true_fn, false_fn, (x,)) rows = torch.export.Dim("rows", min=1, max=8) columns = torch.export.Dim("columns", min=1, max=8) - verify_model_numerically( - CondShapeModel(), - (torch.zeros(2, 3),), + exported_program = export( + CondShapeEqualityModel(), + args=(torch.zeros(3, 3),), + dynamic_shapes={"x": {0: rows, 1: columns}}, + ) + mod = from_exported_program(exported_program) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape, expected_value in (((3, 3), 1.0), ((2, 3), -1.0)): + torch_input = torch.zeros(shape, dtype=torch.float32) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy()))[0] + np.testing.assert_array_equal(actual.numpy(), np.full(shape, expected_value, "float32")) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_cond_shape_inequality_predicate(): + class CondShapeInequalityModel(Module): + def forward(self, x): + def true_fn(x): + return x + 1.0 + + def false_fn(x): + return x - 1.0 + + return torch.cond(x.shape[0] != x.shape[1], true_fn, false_fn, (x,)) + + rows = torch.export.Dim("rows", min=1, max=8) + columns = torch.export.Dim("columns", min=1, max=8) + exported_program = export( + CondShapeInequalityModel(), + args=(torch.zeros(2, 3),), dynamic_shapes={"x": {0: rows, 1: columns}}, - input_sets=[(torch.zeros(shape),) for shape in ((2, 3), (3, 3))], - rtol=0, - atol=0, ) + mod = from_exported_program(exported_program) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape, expected_value in (((2, 3), 1.0), ((3, 3), -1.0)): + torch_input = torch.zeros(shape, dtype=torch.float32) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy()))[0] + np.testing.assert_array_equal(actual.numpy(), np.full(shape, expected_value, "float32")) def test_cond_tuple_output(): @@ -8827,13 +9588,23 @@ class AffineGrid(Module): def forward(self, theta): return torch.nn.functional.affine_grid(theta, [2, 3, 8, 12], align_corners=True) - verify_model_numerically( - AffineGrid(), - (torch.randn(2, 2, 3),), - rtol=1e-5, - atol=1e-5, - run_ep_decomposition=False, - ) + model = AffineGrid() + example_args = (torch.randn(2, 2, 3, dtype=torch.float32),) + + with torch.no_grad(): + pytorch_output = model(*example_args) + + exported_program = export(model, args=example_args) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + + exe = tvm.compile(mod, target="llvm") + vm = relax.VirtualMachine(exe, tvm.cpu()) + + tvm_args = [tvm.runtime.tensor(arg.numpy()) for arg in example_args] + tvm_output = vm["main"](*tvm_args) + tvm_output_np = tvm_output[0].numpy() + + tvm.testing.assert_allclose(tvm_output_np, pytorch_output.numpy(), rtol=1e-5, atol=1e-5) if __name__ == "__main__": diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index 2126032c4fcf..a0ba7971db0a 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -33,10 +33,10 @@ from tvm.script import tirx as T -def verify_model(torch_model, input_info, binding, expected, **import_options): +def verify_model(torch_model, input_info, binding, expected): graph_model = fx.symbolic_trace(torch_model) with torch.no_grad(): - mod = from_fx(graph_model, input_info, **import_options) + mod = from_fx(graph_model, input_info) binding = {k: tvm.runtime.tensor(v) for k, v in binding.items()} expected = relax.transform.BindParams("main", binding)(expected) tvm.ir.assert_structural_equal(mod, expected) @@ -2007,6 +2007,46 @@ def main( binding = {} verify_model(model, input_info, binding, expected2) + class LayerNorm3(Module): + def __init__(self, shape): + super().__init__() + self.shape = shape + self.weight = torch.nn.Parameter(torch.ones(shape)) + self.bias = torch.nn.Parameter(torch.zeros(shape)) + + def forward(self, input): + return torch.nn.functional.layer_norm(input, self.shape, self.weight, self.bias, 1e-5) + + @tvm.script.ir_module + class expected3: + @R.function + def main( + input_1: R.Tensor((1, 3, 10, 10), dtype="float32"), + w1: R.Tensor([10, 10], dtype="float32"), + w2: R.Tensor([10, 10], dtype="float32"), + ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): + # block 0 + with R.dataflow(): + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm( + input_1, + w1, + w2, + axes=[-2, -1], + epsilon=1e-05, + center=True, + scale=True, + ) + gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv + R.output(gv) + return gv + + model = LayerNorm3([10, 10]) + binding = { + "w1": model.weight.detach().numpy(), + "w2": model.bias.detach().numpy(), + } + verify_model(model, input_info, binding, expected3) + def test_cross_entropy(): input_info = [([3, 2], "float32"), ([3], "int32")] @@ -2516,6 +2556,112 @@ def main( verify_model(DivFloorModel(), input_info, {}, expected_div_floor) +def test_round_decimals(): + """torch.round(x, decimals) through from_fx must match PyTorch's round-half-to-even + results, including negative decimals (round(25, -1) == 20). The previous + scale-by-10**decimals implementation multiplied by 0.1 for negative decimals, which + is numerically wrong: 25 * 0.1 == 2.5000000000000004 in float64 rounds up to 30. + """ + input_info = [([10], "float32")] + x = torch.tensor( + [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25], dtype=torch.float32 + ) + + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals + + def forward(self, input): + return torch.round(input, decimals=self.decimals) + + for decimals in (0, 1, -1, -2): + gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval()) + mod = from_fx(gm, input_info) + ex = relax.build(mod, target="llvm") + vm = relax.VirtualMachine(ex, tvm.cpu()) + tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) + got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() + tvm.testing.assert_allclose( + got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6, atol=1e-6 + ) + + +def test_round_decimals_low_precision(): + """Scaling for low-precision inputs must happen in float32 and be cast back. + + 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 == 250000 + exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5 + already overflows float16 (the scale itself becomes inf), turning decimals=5 + and -5 into NaN. Upcasting the input to float32 keeps the scaling exact; the + rounded result is cast back to the input dtype. + """ + input_info = [([8], "float16")] + x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5], dtype=torch.float16) + + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals + + def forward(self, input): + return torch.round(input, decimals=self.decimals) + + # Positive decimals exercise the multiply-by-10**d overflow (4, 5); + # negative decimals exercise the 10**|d| scale overflowing float16 (-5). + for decimals in (2, 4, 5, -2, -4, -5): + gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval()) + mod = from_fx(gm, input_info) + ex = relax.build(mod, target="llvm") + vm = relax.VirtualMachine(ex, tvm.cpu()) + tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) + got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() + tvm.testing.assert_allclose( + got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6, atol=1e-6 + ) + + +def test_round_decimals_large(): + """A large |decimals| must import and run without OverflowError. + + The scale 10**|decimals| used to be built as an unbounded host Python int + before being handed to relax.const, whose int-to-float conversion raises + OverflowError ("int too large to convert to float") once |decimals| >= 309 + (10**309 already exceeds the float64 range). PyTorch accepts such decimals -- + torch.round(x, decimals=309) -- and traces a valid round.decimals call, so + importing the graph must not crash on them. The scale is now built directly + in the float dtype and saturates to inf once it leaves the finite range, + matching PyTorch, whose all-NaN result here comes from the same inf scale. + """ + input_info = [([5], "float32")] + x = torch.tensor([0.5, 1.5, 25.0, -0.5, 0.0], dtype=torch.float32) + + class RoundDecimalsModel(Module): + def __init__(self, decimals): + super().__init__() + self.decimals = decimals + + def forward(self, input): + return torch.round(input, decimals=self.decimals) + + for decimals in (309, -309): + gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval()) + mod = from_fx(gm, input_info) # used to raise OverflowError here + ex = relax.build(mod, target="llvm") + vm = relax.VirtualMachine(ex, tvm.cpu()) + tvm_out = vm["main"](tvm.runtime.tensor(x.numpy())) + got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else tvm_out[0].numpy() + + # The scale overflows to inf, and IEEE arithmetic turns every element into + # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining + # (empty here) finite elements separately, since allclose fails on NaN. + expected = torch.round(x, decimals=decimals) + actual = torch.as_tensor(got) + assert torch.equal(torch.isnan(actual), torch.isnan(expected)) + finite = ~torch.isnan(expected) + assert torch.allclose(actual[finite], expected[finite], rtol=1e-6, atol=1e-6) + + def test_size(): input_info = [([1, 3, 10, 10], "float32")] @@ -3500,20 +3646,37 @@ def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( verify_model(Trunc(), input_info, {}, expected_trunc) -@pytest.mark.parametrize( - "torch_op, relax_op", - [ - (torch.logical_and, R.logical_and), - (torch.logical_or, R.logical_or), - (torch.logical_xor, R.logical_xor), - ], -) -def test_logical_binary(torch_op, relax_op): +def test_logical_and(): + input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")] + + class LogicalAnd(Module): + def forward(self, lhs, rhs): + return torch.logical_and(lhs, rhs) + + @tvm.script.ir_module + class expected: + @R.function + def main( + lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + rhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + ) -> R.Tensor((1, 3, 10, 10), dtype="bool"): + with R.dataflow(): + lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") + lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") + lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_and(lv, lv1) + gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2 + R.output(gv) + return gv + + verify_model(LogicalAnd(), input_info, {}, expected) + + +def test_logical_or(): input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")] - class LogicalBinary(Module): + class LogicalOr(Module): def forward(self, lhs, rhs): - return torch_op(lhs, rhs) + return torch.logical_or(lhs, rhs) @tvm.script.ir_module class expected: @@ -3525,12 +3688,37 @@ def main( with R.dataflow(): lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") - lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(lv, lv1) + lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_or(lv, lv1) gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2 R.output(gv) return gv - verify_model(LogicalBinary(), input_info, {}, expected) + verify_model(LogicalOr(), input_info, {}, expected) + + +def test_logical_xor(): + input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")] + + class LogicalXor(Module): + def forward(self, lhs, rhs): + return torch.logical_xor(lhs, rhs) + + @tvm.script.ir_module + class expected: + @R.function + def main( + lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + rhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + ) -> R.Tensor((1, 3, 10, 10), dtype="bool"): + with R.dataflow(): + lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool") + lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool") + lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_xor(lv, lv1) + gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2 + R.output(gv) + return gv + + verify_model(LogicalXor(), input_info, {}, expected) def test_pow_integer(): @@ -3555,100 +3743,403 @@ def main(inp_0: R.Tensor((4,), dtype="int64")) -> R.Tensor((4,), dtype="int64"): verify_model(Pow(), input_info, {}, expected) -@pytest.mark.parametrize( - "shape, layout, kwargs, size, method, coordinate_mode", - [ - ((1, 3, 10, 10), "NCHW", {"size": (5, 5)}, (5, 5), "nearest_neighbor", "asymmetric"), - ( - (1, 3, 10, 10), - "NCHW", - {"scale_factor": 2.0, "mode": "bilinear", "align_corners": False}, - (20, 20), - "linear", - "half_pixel", - ), - ( - (1, 3, 10, 10), - "NCHW", - {"scale_factor": (2.0, 1.0), "mode": "bicubic", "align_corners": False}, - (20, 10), - "cubic", - "half_pixel", - ), - ( - (1, 3, 4, 10, 10), - "NCDHW", - {"scale_factor": (2.0, 4.0, 4.0), "mode": "trilinear", "align_corners": False}, - (8, 40, 40), - "linear", - "half_pixel", - ), - ( - (1, 3, 4, 10, 10), - "NCDHW", - {"size": (8, 40, 40), "mode": "trilinear", "align_corners": True}, - (8, 40, 40), - "linear", - "align_corners", - ), - ((1, 10, 10, 3), "NHWC", {"size": (5, 5)}, (5, 5), "nearest_neighbor", "asymmetric"), - ( - (1, 10, 10, 3), - "NHWC", - {"scale_factor": 2.0, "mode": "bilinear", "align_corners": False}, - (20, 20), - "linear", - "half_pixel", - ), - ( - (1, 4, 10, 10, 3), - "NDHWC", - {"scale_factor": (2.0, 4.0, 4.0), "mode": "trilinear", "align_corners": True}, - (8, 40, 40), - "linear", - "align_corners", - ), - ], - ids=[ - "nearest", - "scalar-scale", - "tuple-scale-cubic", - "trilinear-scale", - "trilinear-size-aligned", - "nhwc-size", - "nhwc-scale", - "ndhwc-scale-aligned", - ], -) -def test_interpolate(shape, layout, kwargs, size, method, coordinate_mode): +def test_interpolate(): + input_info = [([1, 3, 10, 10], "float32")] + class Interpolate(Module): - def forward(self, x): - return F.interpolate(x, **kwargs) + def forward(self, input): + return torch.nn.functional.interpolate(input, (5, 5)) - resize = R.image.resize3d if len(shape) == 5 else R.image.resize2d - rounding_method = "" if len(shape) == 5 else "round" + @tvm.script.ir_module + class expected1: + @R.function + def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( + (1, 3, 5, 5), dtype="float32" + ): + # block 0 + with R.dataflow(): + lv: R.Tensor((1, 3, 5, 5), dtype="float32") = R.image.resize2d( + input_1, + (5, 5), + roi=[0.000000, 0.000000, 0.000000, 0.000000], + layout="NCHW", + method="nearest_neighbor", + coordinate_transformation_mode="asymmetric", + rounding_method="round", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 3, 5, 5), dtype="float32") = lv + R.output(gv) + return gv - x = relax.Var("x", relax.TensorType(shape, "float32")) - builder = relax.BlockBuilder() - with builder.function("main", [x]): - with builder.dataflow(): - resized = builder.emit( - resize( - x, - size, - layout=layout, - method=method, - coordinate_transformation_mode=coordinate_mode, - rounding_method=rounding_method, + verify_model(Interpolate(), input_info, {}, expected1) + + class Interpolate2(Module): + def forward(self, input): + return torch.nn.functional.interpolate( + input, + size=None, + scale_factor=2.0, + mode="bilinear", + align_corners=False, + ) + + @tvm.script.ir_module + class expected2: + @R.function + def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( + (1, 3, 20, 20), dtype="float32" + ): + # block 0 + with R.dataflow(): + lv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.image.resize2d( + input_1, + (20, 20), + roi=[0.000000, 0.000000, 0.000000, 0.000000], + layout="NCHW", + method="linear", + coordinate_transformation_mode="half_pixel", + rounding_method="round", cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, ) + gv: R.Tensor((1, 3, 20, 20), dtype="float32") = lv + R.output(gv) + return gv + + verify_model(Interpolate2(), input_info, {}, expected2) + + class Interpolate3(Module): + def forward(self, input): + return torch.nn.functional.interpolate( + input, + size=None, + scale_factor=(2.0, 1.0), + mode="bilinear", + align_corners=False, ) - output = builder.emit_output(resized) - builder.emit_func_output(output) - verify_model( - Interpolate(), [(shape, "float32")], {}, builder.get(), default_image_layout=layout - ) + @tvm.script.ir_module + class expected3: + @R.function + def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( + (1, 3, 20, 10), dtype="float32" + ): + # block 0 + with R.dataflow(): + lv: R.Tensor((1, 3, 20, 10), dtype="float32") = R.image.resize2d( + input_1, + (20, 10), + roi=[0.000000, 0.000000, 0.000000, 0.000000], + layout="NCHW", + method="linear", + coordinate_transformation_mode="half_pixel", + rounding_method="round", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 3, 20, 10), dtype="float32") = lv + R.output(gv) + return gv + + verify_model(Interpolate3(), input_info, {}, expected3) + + class Interpolate4(Module): + def forward(self, input): + return torch.nn.functional.interpolate( + input, + size=None, + scale_factor=(2.0, 1.0), + mode="bicubic", + align_corners=False, + ) + + @tvm.script.ir_module + class expected4: + @R.function + def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor( + (1, 3, 20, 10), dtype="float32" + ): + # block 0 + with R.dataflow(): + lv: R.Tensor((1, 3, 20, 10), dtype="float32") = R.image.resize2d( + input_1, + (20, 10), + roi=[0.000000, 0.000000, 0.000000, 0.000000], + layout="NCHW", + method="cubic", + coordinate_transformation_mode="half_pixel", + rounding_method="round", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 3, 20, 10), dtype="float32") = lv + R.output(gv) + return gv + + verify_model(Interpolate4(), input_info, {}, expected4) + + input_info_5d = [([1, 3, 4, 10, 10], "float32")] + + class Interpolate6(Module): + def forward(self, input): + return torch.nn.functional.interpolate( + input, + size=None, + scale_factor=(2.0, 4.0, 4.0), + mode="trilinear", + align_corners=False, + ) + + @tvm.script.ir_module + class expected6: + @R.function + def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor( + (1, 3, 8, 40, 40), dtype="float32" + ): + with R.dataflow(): + lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d( + input_5, + (8, 40, 40), + roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], + layout="NCDHW", + method="linear", + coordinate_transformation_mode="half_pixel", + rounding_method="", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv + R.output(gv) + return gv + + verify_model(Interpolate6(), input_info_5d, {}, expected6) + + class Interpolate7(Module): + def forward(self, input): + return torch.nn.functional.interpolate( + input, + size=(8, 40, 40), + mode="trilinear", + align_corners=False, + ) + + @tvm.script.ir_module + class expected7: + @R.function + def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor( + (1, 3, 8, 40, 40), dtype="float32" + ): + with R.dataflow(): + lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d( + input_5, + (8, 40, 40), + roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], + layout="NCDHW", + method="linear", + coordinate_transformation_mode="half_pixel", + rounding_method="", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv + R.output(gv) + return gv + + verify_model(Interpolate7(), input_info_5d, {}, expected7) + + class Interpolate8(Module): + def forward(self, input): + return torch.nn.functional.interpolate( + input, + size=(8, 40, 40), + mode="trilinear", + align_corners=True, + ) + + @tvm.script.ir_module + class expected8: + @R.function + def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor( + (1, 3, 8, 40, 40), dtype="float32" + ): + with R.dataflow(): + lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d( + input_5, + (8, 40, 40), + roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], + layout="NCDHW", + method="linear", + coordinate_transformation_mode="align_corners", + rounding_method="", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv + R.output(gv) + return gv + + verify_model(Interpolate8(), input_info_5d, {}, expected8) + + +def test_interpolate_nhwc_layout(): + input_info = [([1, 10, 10, 3], "float32")] + + class InterpolateNHWC(Module): + def forward(self, input): + return torch.nn.functional.interpolate(input, (5, 5)) + + @tvm.script.ir_module + class expected_nhwc: + @R.function + def main(input_1: R.Tensor((1, 10, 10, 3), dtype="float32")) -> R.Tensor( + (1, 5, 5, 3), dtype="float32" + ): + # block 0 + with R.dataflow(): + lv: R.Tensor((1, 5, 5, 3), dtype="float32") = R.image.resize2d( + input_1, + (5, 5), + roi=[0.000000, 0.000000, 0.000000, 0.000000], + layout="NHWC", + method="nearest_neighbor", + coordinate_transformation_mode="asymmetric", + rounding_method="round", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 5, 5, 3), dtype="float32") = lv + R.output(gv) + return gv + + # Test with NHWC layout + graph_model = fx.symbolic_trace(InterpolateNHWC()) + with torch.no_grad(): + mod = from_fx(graph_model, input_info, default_image_layout="NHWC") + tvm.ir.assert_structural_equal(mod, expected_nhwc) + + # Test with bilinear interpolation and NHWC layout + class InterpolateNHWC2(Module): + def forward(self, input): + return torch.nn.functional.interpolate( + input, size=None, scale_factor=2.0, mode="bilinear", align_corners=False + ) + + @tvm.script.ir_module + class expected_nhwc2: + @R.function + def main(input_1: R.Tensor((1, 10, 10, 3), dtype="float32")) -> R.Tensor( + (1, 20, 20, 3), dtype="float32" + ): + # block 0 + with R.dataflow(): + lv: R.Tensor((1, 20, 20, 3), dtype="float32") = R.image.resize2d( + input_1, + (20, 20), + roi=[0.000000, 0.000000, 0.000000, 0.000000], + layout="NHWC", + method="linear", + coordinate_transformation_mode="half_pixel", + rounding_method="round", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 20, 20, 3), dtype="float32") = lv + R.output(gv) + return gv + + graph_model2 = fx.symbolic_trace(InterpolateNHWC2()) + with torch.no_grad(): + mod2 = from_fx(graph_model2, input_info, default_image_layout="NHWC") + tvm.ir.assert_structural_equal(mod2, expected_nhwc2) + + input_info_5d = [([1, 4, 10, 10, 3], "float32")] + + class InterpolateNHWC3(Module): + def forward(self, input): + return torch.nn.functional.interpolate( + input, + size=None, + scale_factor=(2.0, 4.0, 4.0), + mode="trilinear", + align_corners=False, + ) + + @tvm.script.ir_module + class expected_nhwc3: + @R.function + def main(input_5: R.Tensor((1, 4, 10, 10, 3), dtype="float32")) -> R.Tensor( + (1, 8, 40, 40, 3), dtype="float32" + ): + with R.dataflow(): + lv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = R.image.resize3d( + input_5, + (8, 40, 40), + roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], + layout="NDHWC", + method="linear", + coordinate_transformation_mode="half_pixel", + rounding_method="", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = lv + R.output(gv) + return gv + + graph_model3 = fx.symbolic_trace(InterpolateNHWC3()) + with torch.no_grad(): + mod3 = from_fx(graph_model3, input_info_5d, default_image_layout="NDHWC") + tvm.ir.assert_structural_equal(mod3, expected_nhwc3) + + class InterpolateNHWC4(Module): + def forward(self, input): + return torch.nn.functional.interpolate( + input, + size=None, + scale_factor=(2.0, 4.0, 4.0), + mode="trilinear", + align_corners=True, + ) + + @tvm.script.ir_module + class expected_nhwc4: + @R.function + def main(input_5: R.Tensor((1, 4, 10, 10, 3), dtype="float32")) -> R.Tensor( + (1, 8, 40, 40, 3), dtype="float32" + ): + with R.dataflow(): + lv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = R.image.resize3d( + input_5, + (8, 40, 40), + roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000], + layout="NDHWC", + method="linear", + coordinate_transformation_mode="align_corners", + rounding_method="", + cubic_alpha=-0.75, + cubic_exclude=0, + extrapolation_value=0, + ) + gv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = lv + R.output(gv) + return gv + + graph_model4 = fx.symbolic_trace(InterpolateNHWC4()) + with torch.no_grad(): + mod4 = from_fx(graph_model4, input_info_5d, default_image_layout="NDHWC") + tvm.ir.assert_structural_equal(mod4, expected_nhwc4) def test_addmm(): @@ -6060,27 +6551,40 @@ def main( R.output(gv) return gv - @I.ir_module - class ExpectedNegative: - @R.function - def main(x: R.Tensor((3, 4), "float32")): - with R.dataflow(): - scaled = R.divide(x, R.const(10.0, "float32")) - rounded = R.round(scaled) - result = R.multiply(rounded, R.const(10.0, "float32")) - output = result - R.output(output) - return output - rounds = [ (0, Expected1), (2, Expected2), - (-1, ExpectedNegative), ] for decimals, expected in rounds: verify_model(Round(decimals), input_info, {}, expected) + # Test numerical accuracy with decimals + test_data = torch.tensor( + [ + [1.2345, 2.3456, 3.4567, 4.5678], + [5.6789, 6.7890, 7.8901, 8.9012], + [9.1234, 10.2345, 11.3456, 12.4567], + ] + ) + + for decimals in [0, 2]: + torch_model = Round(decimals) + graph_model = fx.symbolic_trace(torch_model) + with torch.no_grad(): + mod = from_fx(graph_model, input_info) + + target = tvm.target.Target("llvm") + ex = relax.build(mod, target) + vm = relax.VirtualMachine(ex, tvm.cpu()) + + torch_result = torch_model(test_data).numpy() + tvm_input = tvm.runtime.tensor(test_data.numpy()) + tvm_result = vm["main"](tvm_input).numpy() + + # Use relaxed tolerance due to floating-point precision in decimal operations + tvm.testing.assert_allclose(tvm_result, torch_result, rtol=1e-3, atol=1e-3) + if __name__ == "__main__": tvm.testing.main()