This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 94613dc701 [REFACTOR][TIRx] Replace legacy specialization with
structural mapping (#20787)
94613dc701 is described below
commit 94613dc701c1a4f6b5caa40b573f42dddb0577f8
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 18 13:20:18 2026 -0400
[REFACTOR][TIRx] Replace legacy specialization with structural mapping
(#20787)
---
.../specialize_function_based_on_callsite.cc | 62 ++-
tirx/include/tvm/tirx/function.h | 42 +-
tirx/python/tvm_tirx/tirx/function.py | 60 +--
tirx/src/tirx/ir/function.cc | 7 +-
tirx/src/tirx/ir/specialize.cc | 466 ---------------------
tirx/tests/python/tirx/base/test_specialize.py | 410 ------------------
tirx/tests/python/tirx/ir/__init__.py | 16 +
tirx/tests/python/tirx/ir/test_function.py | 65 +++
8 files changed, 146 insertions(+), 982 deletions(-)
diff --git a/src/relax/transform/specialize_function_based_on_callsite.cc
b/src/relax/transform/specialize_function_based_on_callsite.cc
index 457a3d72ce..4c7776769c 100644
--- a/src/relax/transform/specialize_function_based_on_callsite.cc
+++ b/src/relax/transform/specialize_function_based_on_callsite.cc
@@ -21,6 +21,7 @@
* \brief Update tirx::Function tensors based on updated scope (or structure)
info.
*/
+#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/nested_msg.h>
@@ -28,6 +29,7 @@
#include <tvm/relax/op_attr.h>
#include <tvm/relax/transform.h>
#include <tvm/tirx/index_map.h>
+#include <tvm/tirx/op/memory.h>
#include <tuple>
@@ -82,7 +84,7 @@ class SpecializeTIRCallArgs : ExprMutator {
auto gv = call->args[0].as_or_throw<GlobalVar>();
auto pfunc = mod_->Lookup(gv).as_or_throw<tirx::Function>();
auto args = call->args[1].as_or_throw<Tuple>()->fields;
- ffi::Map<tvm::Var, ffi::Variant<TensorVar, Expr>> param_map;
+ ffi::Map<tvm::Var, TensorVar> param_map;
for (size_t i = 0; i < args.size(); ++i) {
auto ty = GetType(args[i]);
@@ -131,6 +133,7 @@ class SpecializeTIRCallArgs : ExprMutator {
"output structinfo, but got "
<< si;
auto ty = si.as_or_throw<TensorType>();
+ scope = runtime::MemoryScope();
if (ty->vdevice.has_value()) {
scope = ty->vdevice.value()->memory_scope;
}
@@ -143,7 +146,62 @@ class SpecializeTIRCallArgs : ExprMutator {
}
}
- auto new_pfunc = Specialize(pfunc, param_map);
+ // call_tir supplies dense tensor descriptors. Bind their symbolic
dimensions
+ // explicitly, retaining each parameter's access contract and storage
scope.
+ sym::Analyzer analyzer;
+ ffi::Map<tvm::Var, Expr> replacements;
+ auto bind = [&](const PrimExpr& symbolic, const PrimExpr& concrete) {
+ if (analyzer->CanProveEqual(symbolic, concrete)) return;
+ auto var = symbolic.as<PrimVar>();
+ TVM_FFI_CHECK(var.has_value() && !replacements.count(var.value()),
ValueError)
+ << "Inconsistent call_tir tensor dimension: " << symbolic << " vs. "
<< concrete;
+ replacements.Set(var.value(), concrete);
+ analyzer->Bind(var.value(), concrete);
+ };
+ for (const auto& param : pfunc->params) {
+ auto descriptor = param_map.Get(param);
+ if (!descriptor) continue;
+ auto tensor = param.as_or_throw<TensorVar>();
+ auto actual = descriptor.value();
+ TVM_FFI_CHECK_EQ(tensor.shape().size(), actual.shape().size(),
ValueError);
+ TVM_FFI_CHECK(tensor.ty()->dtype == actual.ty()->dtype, ValueError);
+ for (size_t i = 0; i < tensor.shape().size(); ++i) {
+ bind(tensor.shape()[i], actual.shape()[i]);
+ }
+ auto source = tirx::MatchShapeStridePattern(tensor.ty());
+ auto target = tirx::MatchShapeStridePattern(actual.ty());
+ TVM_FFI_CHECK(source.has_value() && target.has_value(), ValueError)
+ << "call_tir requires shape-aligned strided tensor layouts";
+ for (size_t i = 0; i < source.value().get<1>().size(); ++i) {
+ bind(source.value().get<1>()[i], target.value().get<1>()[i]);
+ }
+ bind(source.value().get<2>(), target.value().get<2>());
+ // Preserve layout and alignment; the Relax descriptor only chooses
scope.
+ auto ty = tensor.ty();
+ replacements.Set(param, TensorVar(param->name,
+ tirx::TensorType(ty->shape, ty->dtype,
actual.scope(),
+ ty->data_alignment,
ty->layout, ty->loc),
+ param->loc));
+ }
+ ffi::Array<tvm::Var> params;
+ for (const auto& param : pfunc->params) {
+ if (!param.as<PrimVar>() || !replacements.count(param))
params.push_back(param);
+ }
+ auto signature = tirx::Function(params, pfunc->body, pfunc->ret_type,
pfunc->attrs, pfunc->loc);
+ auto mapped = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
+ signature,
+ [&](const tvm::Var& var) -> Expr { return
replacements.Get(var).value_or(var); });
+ auto new_pfunc = ffi::StructuralMap<ffi::WalkOrder::kPostOrder>(
+ mapped,
+ [](const tvm::Call& call) {
+ if (!call->op.same_as(tirx::tensor_data_ptr_op()))
return call;
+ // Pointer projections inherit the newly bound
tensor storage scope.
+ auto copy =
ffi::make_object<tvm::CallNode>(*call.get());
+ copy->ty = tvm::Call::ReinferType(copy.get());
+ return tvm::Call(std::move(copy));
+ },
+ [&](const PrimExpr& expr) { return
analyzer->Simplify(expr); })
+ .as_or_throw<tirx::Function>();
auto new_function =
WithAttr(new_pfunc, tvm::relax::attr::kScoped.name,
static_cast<int64_t>(1));
updates_->Add(gv, new_function);
diff --git a/tirx/include/tvm/tirx/function.h b/tirx/include/tvm/tirx/function.h
index 411412c23f..75d751276f 100644
--- a/tirx/include/tvm/tirx/function.h
+++ b/tirx/include/tvm/tirx/function.h
@@ -25,7 +25,6 @@
#define TVM_TIRX_FUNCTION_H_
#include <tvm/ffi/container/map.h>
-#include <tvm/ffi/container/variant.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ir/cow.h>
@@ -70,9 +69,7 @@ class FunctionNode : public BaseFuncNode {
bool SEqual(const FunctionNode* other,
ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal)
const {
- // `ty` is derived from the fields below. Function transformations update
- // those source fields without maintaining this redundant cache eagerly.
- // Remove this exception once all Function mutation paths recompute `ty`.
+ // `ty` is derived from the signature; compare its source fields directly.
return equal(attrs, other->attrs, false, "attrs") &&
equal(params, other->params, true, "params") &&
equal(ret_type, other->ret_type, false, "ret_type") &&
@@ -140,43 +137,6 @@ class Function : public BaseFunc {
*/
TVM_TIRX_DLL Function RenewDef(Function func);
-/*!
- * \brief Specialize parameters of Function.
- * \param func The Function to be specialized.
- * \param param_map The mapping from function params to the instance.
- * \return The new function with parameter specialized.
- * \note We can define a TIRx function with symbolic shape:
- *
- * \code{.py}
- * from __future__ import annotations
- *
- * @T.function
- * def mem_copy(A: T.Tensor((m, n), "float32"), B: T.Tensor((m, n),
"float32"),
- * m: T.int32, n: T.int32) -> None:
- * for i, j in T.grid(m, n):
- * B[i, j] = A[i, j]
- * \endcode
- *
- * Then we can make it specialized with given shapes or tensors.
- *
- * \code{.py}
- * a, _, m, n = mem_copy.params
- * func = mem_copy.specialize({a: tirx.decl_tensor((16, 16))})
- * # or
- * func = mem_copy.specialize({n: 16, m: 16})
- * \endcode
- *
- * \code{.py}
- * @T.function
- * def mem_copy_16_16(A: T.Tensor((16, 16), "float32"),
- * B: T.Tensor((16, 16), "float32")) -> None:
- * for i, j in T.grid(16, 16):
- * B[i, j] = A[i, j]
- * \endcode
- */
-TVM_TIRX_DLL Function Specialize(Function func,
- const ffi::Map<Var, ffi::Variant<TensorVar,
Expr>>& param_map);
-
/*!
* \brief Function specific attribute names.
*
diff --git a/tirx/python/tvm_tirx/tirx/function.py
b/tirx/python/tvm_tirx/tirx/function.py
index 9ce3bd1487..8b06d4678d 100644
--- a/tirx/python/tvm_tirx/tirx/function.py
+++ b/tirx/python/tvm_tirx/tirx/function.py
@@ -17,7 +17,7 @@
# pylint: disable=unrecognized-inline-option
"""Function data types."""
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
import tvm_ffi
@@ -27,7 +27,6 @@ from tvm_tirx.ir import BaseFunc, Scriptable, SeqStmt, Stmt,
Var
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from . import _ffi_api
-from .expr import Expr
@tvm_ffi.register_object("tirx.Function")
@@ -107,63 +106,6 @@ class Function(BaseFunc, Scriptable):
loc=loc,
)
- def specialize(self, param_map: Mapping[Var, Expr | Var]):
- """Specialize parameters of Function
-
- Parameters
- ----------
-
- param_map : Mapping[Var, Union[Expr, Var]]
- The mapping from function params to the instance
-
- Examples
- --------
- We can define a TIRX function with symbolic shape:
-
- .. code-block:: python
-
- from __future__ import annotations
-
- @T.function
- def mem_copy(
- A: T.Tensor((m, n), "float32"),
- B: T.Tensor((m, n), "float32"),
- m: T.int32,
- n: T.int32,
- ) -> None:
-
- for i, j in T.grid(m, n):
- B[i, j] = A[i, j]
-
- Then we can make it specialized with given shapes or tensors.
-
- .. code-block:: python
-
- a, _, m, n = mem_copy.params
- tensor = tirx.Var("tensor", tirx.TensorType((16, 16), "float32"))
- func = mem_copy.specialize({a: tensor})
- # or
- func = mem_copy.specialize({n: 16, m: 16})
-
- The specialized function:
-
- .. code-block:: python
-
- @T.function
- def mem_copy_16_16(
- A: T.Tensor((16, 16), "float32"), B: T.Tensor((16, 16),
"float32")
- ) -> None:
-
- for i, j in T.grid(16, 16):
- B[i, j] = A[i, j]
-
- Returns
- -------
- func : Function
- The new function with parameter specialized
- """
- return _ffi_api.Specialize(self, param_map) # type: ignore
-
def renew_def(func: Function) -> Function:
"""Renew variable definitions and remap their uses in a function.
diff --git a/tirx/src/tirx/ir/function.cc b/tirx/src/tirx/ir/function.cc
index 2e425aff1e..e7b225d9c1 100644
--- a/tirx/src/tirx/ir/function.cc
+++ b/tirx/src/tirx/ir/function.cc
@@ -70,6 +70,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
FunctionMutate(
?
ffi::Optional<SeqStmt>(SeqStmt(std::move(replacement).value()))
: std::nullopt;
}
+ copy->ty = copy->func_type_annotation();
return ffi::Any(std::move(copy));
}
@@ -100,6 +101,8 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
FunctionMaybeInplaceMut
?
ffi::Optional<SeqStmt>(SeqStmt(std::move(replacement).value()))
: std::nullopt;
}
+ // A nested signature node may change in place without replacing params or
ret_type.
+ self->ty = self->func_type_annotation();
return ffi::Unchanged();
}
@@ -143,10 +146,6 @@ Function RenewDef(Function func) {
Var fresh(var->name, var->ty, var->loc);
remap.Set(var, fresh);
return fresh;
- },
- [](const Function& mapped) {
- return Function(mapped->params, mapped->body, mapped->ret_type,
mapped->attrs,
- mapped->loc);
})
.as_or_throw<Function>();
}
diff --git a/tirx/src/tirx/ir/specialize.cc b/tirx/src/tirx/ir/specialize.cc
deleted file mode 100644
index 7abf56b4e3..0000000000
--- a/tirx/src/tirx/ir/specialize.cc
+++ /dev/null
@@ -1,466 +0,0 @@
-/*
- * 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 src/tirx/ir/specialize.cc
- * \brief Specialize parameters of Function.
- */
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/extra/structural_equal.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/tirx/function.h>
-#include <tvm/tirx/layout.h>
-#include <tvm/tirx/op/memory.h>
-#include <tvm/tirx/stmt_functor.h>
-
-#include <functional>
-#include <unordered_set>
-
-#include "../transform/ir_utils.h"
-
-namespace tvm {
-namespace tirx {
-
-using VarMap = std::unordered_map<Var, Expr>;
-
-/**************** Helper functions ****************/
-
-/*! \brief Helper function to check whether the given var is in function
parameter list. */
-inline bool IsParam(const Function& func, const Var& param) {
- return std::any_of(func->params.begin(), func->params.end(),
- [&](const Var& var) { return var.same_as(param); });
-}
-
-/**************** Specializer ****************/
-
-// Try fold constants if op's child get specialized to constant.
-#define DEFINE_SPECIALIZER_BINARY_OP_MUTATE(BinaryNode, BinaryFunc)
\
- UnchangedOr<PrimExpr> Mutate_(const BinaryNode* op, InplaceMode
inplace_mode) final { \
- auto a_result = Mutate(op->a, inplace_mode);
\
- bool a_unchanged = a_result.UnchangedOrSameAs(op->a);
\
- PrimExpr a = std::move(a_result).ValueOrUnchanged(op->a);
\
- auto b_result = Mutate(op->b, inplace_mode);
\
- bool b_unchanged = b_result.UnchangedOrSameAs(op->b);
\
- PrimExpr b = std::move(b_result).ValueOrUnchanged(op->b);
\
- if (a_unchanged && b_unchanged) {
\
- return ffi::Unchanged();
\
- } else {
\
- return BinaryFunc(a, b, op->loc);
\
- }
\
- }
-#define DEFINE_SPECIALIZER_UNARY_OP_MUTATE(UnaryNode, UnaryFunc)
\
- UnchangedOr<PrimExpr> Mutate_(const UnaryNode* op, InplaceMode inplace_mode)
final { \
- auto a_result = Mutate(op->a, inplace_mode);
\
- bool a_unchanged = a_result.UnchangedOrSameAs(op->a);
\
- PrimExpr a = std::move(a_result).ValueOrUnchanged(op->a);
\
- if (a_unchanged) {
\
- return ffi::Unchanged();
\
- } else {
\
- return UnaryFunc(a, op->loc);
\
- }
\
- }
-
-/*! \brief Mutator to specialize function and remove const parameters */
-class FunctionSpecializer : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- explicit FunctionSpecializer(const VarMap& var_map) {
- for (const auto& [var, value] : var_map) {
- if (!var.as<TensorVar>()) VarRemapSet(var, value);
- }
- }
-
- static Function Specialize(Function f, const VarMap& var_map) {
- auto specializer = ffi::make_object<FunctionSpecializer>(var_map);
- for (const Var& param : f->params) {
- auto tensor = param.as<TensorVar>();
- auto replacement = var_map.find(param);
- if (!tensor || replacement == var_map.end()) {
- continue;
- }
- if (auto replacement_var = replacement->second.as<Var>()) {
- if (auto replacement_tensor = replacement_var.value().as<TensorVar>())
{
- if (IsParam(f, replacement_var.value())) {
- specializer->VarRemapSet(tensor.value(),
replacement_tensor.value());
- } else {
- specializer->constrained_tensor_params_.insert(param.get());
- specializer->tensor_storage_scopes_.emplace(tensor.value(),
-
replacement_tensor.value().scope());
- }
- }
- }
- }
-
- // Updating parameters
- ffi::Array<Var> params;
- bool param_updated = false;
- for (const auto& var : f->params) {
- Var new_var = var;
- if (auto tensor = var.as<TensorVar>()) {
- TensorVar new_tensor = specializer->MutateTensor(tensor.value());
- new_var = new_tensor;
- if (!new_tensor.same_as(tensor.value())) {
- param_updated = true;
- specializer->VarRemapSet(tensor.value(), new_tensor);
- specializer->defined_tensors_.insert(tensor.value().get());
- }
- }
- // Remove parmeters which has been specialized.
- if (var_map.find(var) == var_map.end() ||
- specializer->constrained_tensor_params_.count(var.get())) {
- params.push_back(new_var);
- } else {
- param_updated = true;
- }
- }
-
- auto planner = ffi::make_object<TensorPlanner>(specializer.get());
- planner->Visit(f->body);
-
- auto body_result =
- specializer->Mutate(f->body, f.unique() ? InplaceMode::kAllow :
InplaceMode::kDisallow);
- bool body_unchanged = body_result.UnchangedOrSameAs(f->body);
- auto body = std::move(body_result).ValueOrUnchanged(f->body);
-
- if (param_updated || !body_unchanged) {
- return Function(params, body, f->ret_type, f->attrs, f->loc);
- } else {
- return f;
- }
- }
-
- private:
- class TensorPlanner : public StmtExprVisitor {
- public:
- using StmtExprVisitor::Visit_;
-
- explicit TensorPlanner(FunctionSpecializer* specializer) :
specializer_(specializer) {}
-
- private:
- ffi::Optional<VisitInterrupt> Visit_(const VarNode* op) final {
- if (op->ty.as<TensorTypeNode>()) {
- if (def_region_kind() == kTVMFFIDefRegionKindSimple) {
- const TensorVar tensor =
ffi::GetRef<Var>(op).as_or_throw<TensorVar>();
- specializer_->MutateAllocTensor(tensor);
- } else {
-
specializer_->ValidateTensorUse(ffi::GetRef<Var>(op).as_or_throw<TensorVar>());
- }
- }
- return StmtExprVisitor::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const BindNode* op) final {
- if (const auto* call = op->value.as<CallNode>();
- call &&
- (call->op.same_as(tirx::alloc_tensor_op()) ||
call->op.same_as(tirx::decl_tensor_op()))) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->WithDefRegionKind(
- kTVMFFIDefRegionKindSimple, [&]() { return this->Visit(op->var);
}));
- // The definition above owns traversal of the complete tensor type.
- // Value operands include declaration data or allocation placement.
- return Visit(call->args);
- }
- return StmtExprVisitor::Visit_(op);
- }
-
- FunctionSpecializer* specializer_;
- };
-
- UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
final {
- auto result = StmtExprMutator::Mutate_(op, inplace_mode);
- if (!result.IsUnchanged()) {
- op = ffi::AnyView(result).as<CallNode>();
- if (!op->unique()) inplace_mode = InplaceMode::kDisallow;
- }
- if (!op->op.same_as(tirx::tensor_data_ptr_op()) || op->args.size() != 1)
return result;
- PtrType type = op->args[0].as_or_throw<TensorVar>().ty()->DataPtrType();
- if (ffi::StructuralEqual()(op->ty, type)) return result;
- if (inplace_mode == InplaceMode::kAllow) {
- const_cast<CallNode*>(op)->ty = std::move(type);
- return result;
- }
- auto copy = ffi::make_object<CallNode>(*op);
- copy->ty = std::move(type);
- return Expr(std::move(copy));
- }
-
- UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode
inplace_mode) final {
- auto result = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow);
- if (result.UnchangedOrSameAs(ffi::GetRef<PrimExpr>(op))) return
ffi::Unchanged();
- auto load = std::move(result).ValueUnchecked().as_or_throw<TensorLoad>();
- if (auto tensor = load->source.as<TensorVar>()) {
- return TensorLoad(tensor.value(), load->indices, load->loc);
- }
- return load;
- }
-
- UnchangedOr<Expr> Mutate_(const TensorRegionNode* op, InplaceMode
inplace_mode) final {
- auto result = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow);
- if (result.UnchangedOrSameAs(ffi::GetRef<Expr>(op))) return
ffi::Unchanged();
- auto region =
std::move(result).ValueUnchecked().as_or_throw<TensorRegion>();
- if (auto tensor = region->source.as<TensorVar>()) {
- return TensorRegion(tensor.value(), region->region, TensorRegionType(),
region->loc);
- }
- return region;
- }
-
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::AddNode, add);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::SubNode, sub);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::MulNode, mul);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::DivNode, div);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::ModNode, truncmod);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::FloorDivNode, floordiv);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::FloorModNode, floormod);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::MaxNode, max);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::MinNode, min);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::EQNode, equal);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::NENode, not_equal);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::LTNode, less);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::LENode, less_equal);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::GTNode, greater);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::GENode, greater_equal);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::AndNode, logical_and);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::OrNode, logical_or);
- DEFINE_SPECIALIZER_UNARY_OP_MUTATE(prim::NotNode, logical_not);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::LShiftNode, left_shift);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::RShiftNode, right_shift);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::BitwiseAndNode, bitwise_and);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::BitwiseOrNode, bitwise_or);
- DEFINE_SPECIALIZER_BINARY_OP_MUTATE(prim::BitwiseXorNode, bitwise_xor);
- DEFINE_SPECIALIZER_UNARY_OP_MUTATE(prim::BitwiseNotNode, prim::BitwiseNot);
- TensorVar MutateTensor(const TensorVar& tensor) {
- ffi::Any mapped = VarRemapGet(tensor);
- if (mapped.type_index() != ffi::TypeIndex::kTVMFFINone) {
- return
std::move(mapped).as_or_throw<UnchangedOr<TensorVar>>().ValueOrUnchanged(tensor);
- }
-
- ffi::Optional<runtime::MemoryScope> specialized_storage_scope;
- if (auto it = tensor_storage_scopes_.find(tensor); it !=
tensor_storage_scopes_.end()) {
- specialized_storage_scope = it->second;
- }
-
- ffi::Array<PrimExpr> shape =
- tensor.shape().Map([this](const PrimExpr& e) { return
Mutate(e).ValueOrUnchanged(e); });
- ffi::Optional<Layout> layout = tensor.ty()->layout;
- auto mapped_layout = Mutate(ffi::AnyView(layout), InplaceMode::kDisallow);
- bool layout_changed = !mapped_layout.UnchangedOrSameAs(layout);
- layout =
std::move(mapped_layout).ValueOrUnchanged(layout).cast<ffi::Optional<Layout>>();
-
- bool storage_scope_changed = specialized_storage_scope.has_value() &&
- specialized_storage_scope.value() !=
tensor.scope();
- if (tensor.shape().same_as(shape) && !layout_changed &&
!storage_scope_changed) {
- return tensor;
- } else {
- auto n = ffi::make_object<TensorTypeNode>(*tensor.ty().get());
- n->shape = std::move(shape);
- if (layout_changed) {
- n->layout = std::move(layout);
- }
- if (storage_scope_changed) {
- n->scope = specialized_storage_scope.value();
- }
- return TensorVar(tensor->name, TensorType(std::move(n)), tensor->loc);
- }
- }
-
- void MutateAllocTensor(const TensorVar& alloc_buf) {
- TVM_FFI_ICHECK(defined_tensors_.insert(alloc_buf.get()).second)
- << "Multiple points of definition found for tensor " << alloc_buf;
- VarRemapSet(alloc_buf, MutateTensor(alloc_buf));
- }
-
- void ValidateTensorUse(const TensorVar& old_tensor) {
- if (VarRemapGet(old_tensor).type_index() != ffi::TypeIndex::kTVMFFINone)
return;
-
- auto mutated = MutateTensor(old_tensor);
- TVM_FFI_ICHECK(mutated.same_as(old_tensor))
- << "TensorVar " << old_tensor << " (shape = " << old_tensor.shape() <<
")"
- << " was used without a declaration, "
- << "and would be specialized into " << mutated << " (shape = " <<
mutated.shape() << "). "
- << "While usage of an undeclared tensor is currently allowed in TIR, "
- << "mutation must occur at the tensor's point of definition "
- << "(see discussion on https://github.com/apache/tvm/pull/14565 for
more details). "
- << "Please add a definition for this tensor, "
- << "either as a TensorType-annotated Function parameter, "
- << "in a block's tensor allocations, "
- << "or in a DeclTensor statement.";
- }
-
- /*! \brief Definition identities used only to validate declaration order. */
- std::unordered_set<const VarNode*> defined_tensors_;
- /*! \brief Storage constraints supplied by concrete, non-parameter tensors.
*/
- std::unordered_map<TensorVar, runtime::MemoryScope, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>
- tensor_storage_scopes_;
- /*! \brief Tensor parameters constrained by a concrete, non-parameter
tensor. */
- std::unordered_set<const VarNode*> constrained_tensor_params_;
-};
-
-/*!
- * \brief Update Specialize var map with tensor matching.
- * \param func The function to be specialized.
- * \param param The given function parameter
- * \param specific_buf The matching tensor.
- * \param var_map The var mapping to be updated.
- * \note This function will match target tensor's shape, strides and
element_offset
- * For example, we define a tensor in Function:
- * A: T.Tensor([m, n])
- *
- * Then we match it with a tensor B = Var("tensor", TensorType((8, 16),
"float32"))
- *
- * It means we have two var mappings here: m = 8 and n = 16
- *
- * If the tensor signature is not a Var, the mapping will fail.
- * e.g. A: T.Tensor([m * 2, n + 1])
- */
-void UpdateSpecializeVarMap(const Function& func, const Var& param, const
TensorVar& specific_buf,
- VarMap* var_map) {
- // preliminaries
- prim::ExprDeepEqual equal;
-
- auto opt_tensor = param.as<TensorVar>();
- TVM_FFI_CHECK(opt_tensor, ValueError)
- << "specialize expects param to have a TensorType annotation";
- const TensorVar& buf_to_specialize = opt_tensor.value();
-
- // build var mapping using specific_buf's parameters
- auto build_var_mapping = [&](const Expr& new_expr, const Expr& old_expr) {
- auto expr_equal = [&](const Expr& lhs, const Expr& rhs) {
- if (auto lhs_prim = lhs.as<PrimExpr>()) {
- auto rhs_prim = rhs.as<PrimExpr>();
- return rhs_prim && equal(lhs_prim.value(), rhs_prim.value());
- }
- return lhs.same_as(rhs);
- };
- Expr resolved = old_expr;
- if (auto prim_expr = old_expr.as<PrimExpr>()) {
- auto mutator = ffi::make_object<FunctionSpecializer>(*var_map);
- resolved =
mutator->Mutate(prim_expr.value()).ValueOrUnchanged(prim_expr.value());
- }
- if (!expr_equal(new_expr, resolved)) {
- auto maybe_var = old_expr.as<Var>();
- TVM_FFI_CHECK(maybe_var, TypeError)
- << "The signature of target tensor exprected an independent Var, but
got " << old_expr
- << ".";
- const Var& var = maybe_var.value();
- auto it = var_map->find(var);
- if (it != var_map->end()) {
- TVM_FFI_CHECK(expr_equal(it->second, new_expr), ValueError)
- << "The assigned value of var " << var << " mismatched. " <<
it->second << " vs. "
- << new_expr << ".";
- } else {
- var_map->insert_or_assign(var, new_expr);
- }
- }
- };
-
- // Check tensor dimensions
- TVM_FFI_CHECK(specific_buf.shape().size() ==
buf_to_specialize.shape().size(), ValueError)
- << "The tensor dimensions mismatched" <<
buf_to_specialize.shape().size() << " vs. "
- << specific_buf.shape().size() << ".";
-
- for (size_t i = 0; i < specific_buf.shape().size(); ++i) {
- build_var_mapping(specific_buf.shape()[i], buf_to_specialize.shape()[i]);
- }
- auto specific = MatchShapeStridePattern(specific_buf.ty());
- auto original = MatchShapeStridePattern(buf_to_specialize.ty());
- TVM_FFI_CHECK(specific.has_value() && original.has_value(), ValueError)
- << "Tensor specialization requires shape-aligned strided memory layouts";
- auto new_strides = specific.value().get<1>();
- auto old_strides = original.value().get<1>();
- for (size_t i = 0; i < old_strides.size(); ++i) {
- build_var_mapping(new_strides[i], old_strides[i]);
- }
- build_var_mapping(specific.value().get<2>(), original.value().get<2>());
- // The specializer distinguishes a concrete tensor constraint (which keeps
- // the parameter) from another function parameter (which aliases it).
- build_var_mapping(specific_buf, buf_to_specialize);
-
- // Check data_alignment.
- // These two signatures are int, so we do not need map them.
- TVM_FFI_CHECK_EQ(specific_buf.ty()->data_alignment,
buf_to_specialize.ty()->data_alignment,
- ValueError)
- << "The tensor data_alignment mismatched" <<
buf_to_specialize.ty()->data_alignment << " vs. "
- << specific_buf.ty()->data_alignment << ".";
-}
-
-/*!
- * \brief Update Specialize var map with parameter value.
- * \param func The function to be specialized.
- * \param param The given function parameter
- * \param specific_expr The parameter value.
- * \param var_map The var mapping to be updated.
- */
-void UpdateSpecializeVarMap(const Function& func, const Var& param, const
Expr& specific_expr,
- VarMap* var_map) {
- // check param is in Function's parameters
- TVM_FFI_CHECK(IsParam(func, param), ValueError)
- << "Specialize expects param to be in Function's params";
- // Specialize a scalar parameter rather than a tensor parameter.
- TVM_FFI_CHECK(!param.as<TensorVar>(), ValueError)
- << "Specialize expects param to not have a TensorType annotation";
- // build var mapping using specific_expr
- var_map->insert_or_assign(param, specific_expr);
-}
-
-/**************** Implementation ****************/
-
-Function Specialize(Function func, const ffi::Map<Var, ffi::Variant<TensorVar,
Expr>>& param_map) {
- VarMap var_map;
- for (const auto& kv : param_map) {
- const Var& param = kv.first;
- const ffi::Variant<TensorVar, Expr>& instance = kv.second;
- if (auto opt_tensor = instance.as<TensorVar>()) {
- UpdateSpecializeVarMap(func, param, opt_tensor.value(), &var_map);
- } else if (auto opt_expr = instance.as<Expr>()) {
- UpdateSpecializeVarMap(func, param, opt_expr.value(), &var_map);
- } else {
- TVM_FFI_THROW(TypeError) << "specialize expected instance to be
TensorVar or Expr";
- }
- }
- return FunctionSpecializer::Specialize(func, std::move(var_map));
-}
-
-/**************** FFI ****************/
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.Specialize", [](Function func,
- const ffi::Map<Var, ffi::Any>&
param_map) {
- VarMap var_map;
- for (const auto& [param, instance] : param_map) {
- if (auto tensor = instance.as<TensorVar>()) {
- UpdateSpecializeVarMap(func, param, tensor.value(), &var_map);
- } else if (const ExprNode* expr = instance.as<ExprNode>()) {
- UpdateSpecializeVarMap(func, param, ffi::GetRef<Expr>(expr), &var_map);
- } else if (instance.type_index() < ffi::TypeIndex::kTVMFFISmallStr) {
- UpdateSpecializeVarMap(func, param, instance.cast<PrimExpr>(),
&var_map);
- } else {
- TVM_FFI_THROW(TypeError) << "specialize expected instance to be
TensorVar or Expr, but got "
- << instance.GetTypeKey();
- }
- }
- return FunctionSpecializer::Specialize(func, std::move(var_map));
- });
-}
-
-} // namespace tirx
-} // namespace tvm
diff --git a/tirx/tests/python/tirx/base/test_specialize.py
b/tirx/tests/python/tirx/base/test_specialize.py
deleted file mode 100644
index 8bdc8685b2..0000000000
--- a/tirx/tests/python/tirx/base/test_specialize.py
+++ /dev/null
@@ -1,410 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=missing-function-docstring, missing-module-docstring
-# ruff: noqa: F401
-
-from __future__ import annotations
-
-import pytest
-
-import tvm_tirx
-from tvm_tirx.script import tirx as T
-
-m = T.dynamic("m", "int32")
-
-
-def assert_structural_equal_ignore_global_symbol(lhs, rhs):
- tvm_tirx.ir.assert_structural_equal(
- lhs.without_attr("global_symbol"), rhs.without_attr("global_symbol")
- )
-
-
[email protected]
-def matmul(
- A: T.Tensor([m, n], "float32"),
- B: T.Tensor([m, n], "float32"),
- C: T.Tensor([m, m], "float32"),
- n: T.int32,
-) -> None:
- for i, j, k in T.grid(m, m, n):
- if k == 0:
- C[i, j] = 0.0
- C[i, j] = C[i, j] + A[i, k] * B[j, k]
-
-
[email protected]
-def matmul_128(
- A: T.Tensor([128, 128], "float32"),
- B: T.Tensor([128, 128], "float32"),
- C: T.Tensor([128, 128], "float32"),
-) -> None:
- for i, j, k in T.grid(128, 128, 128):
- if k == 0:
- C[i, j] = 0.0
- C[i, j] = C[i, j] + A[i, k] * B[j, k]
-
-
-m = T.dynamic("m", "int32")
-
-
[email protected]
-def matmul_m_128(
- A: T.Tensor([m, 128], "float32"),
- B: T.Tensor([m, 128], "float32"),
- C: T.Tensor([m, m], "float32"),
-) -> None:
- for i, j, k in T.grid(m, m, 128):
- if k == 0:
- C[i, j] = 0.0
- C[i, j] = C[i, j] + A[i, k] * B[j, k]
-
-
-# x is considered undefined because it appears as part of x*8,
-# but not on its own
-x = T.dynamic("x", "int32")
-m = T.dynamic("m", "int32")
-
-
[email protected](check_well_formed=False)
-def matmul_m_8x(
- A: T.Tensor([m, x * 8], "float32"),
- B: T.Tensor([m, x * 8], "float32"),
- C: T.Tensor([m, m], "float32"),
-) -> None:
- for i, j, k in T.grid(m, m, x * 8):
- if k == 0:
- C[i, j] = 0.0
- C[i, j] = C[i, j] + A[i, k] * B[j, k]
-
-
-m = T.dynamic("m", "int32")
-n = T.dynamic("n", "int32")
-
-
[email protected]
-def element_wise(A: T.Tensor((m, n), "float32"), C: T.Tensor((m, n),
"float32")) -> None:
- B = T.alloc_tensor(ty_args=[T.Tensor((m, n), "float32")])
-
- for i, j in T.grid(m, n):
- B[i, j] = A[i, j] * 2.0
-
- for i, j in T.grid(m, n):
- C[i, j] = B[i, j] + 1.0
-
-
[email protected]
-def element_wise_128_64(
- A: T.Tensor((128, 64), "float32"), C: T.Tensor((128, 64), "float32")
-) -> None:
- B = T.alloc_tensor(ty_args=[T.Tensor((128, 64), "float32")])
-
- for i, j in T.grid(128, 64):
- B[i, j] = A[i, j] * 2.0
-
- for i, j in T.grid(128, 64):
- C[i, j] = B[i, j] + 1.0
-
-
-n = T.dynamic("n", "int32")
-
-
[email protected]
-def element_wise_128_n(A: T.Tensor((128, n), "float32"), C: T.Tensor((128, n),
"float32")) -> None:
- B = T.alloc_tensor(ty_args=[T.Tensor((128, n), "float32")])
-
- for i, j in T.grid(128, n):
- B[i, j] = A[i, j] * 2.0
-
- for i, j in T.grid(128, n):
- C[i, j] = B[i, j] + 1.0
-
-
-mem_copy_m = T.dynamic("mem_copy_m", "int32")
-
-mem_copy_n = T.dynamic("mem_copy_n", "int32")
-
-
[email protected]
-def mem_copy(
- A: T.Tensor(
- (mem_copy_m, mem_copy_n),
- "float32",
- layout=T.TileLayout(T.S[(mem_copy_m, mem_copy_n) : [p, 1]] + (q)), #
noqa: F821
- ),
- B: T.Tensor(
- (mem_copy_m, mem_copy_n),
- "float32",
- layout=T.TileLayout(T.S[(mem_copy_m, mem_copy_n) : [p, 1]] + (q)), #
noqa: F821
- ),
- m: mem_copy_m,
- n: mem_copy_n,
- p: T.int32,
- q: T.int32,
-) -> None:
- for i, j in T.grid(m, n):
- B[i, j] = A[i, j]
-
-
[email protected]
-def mem_copy_16_16_8_4(
- A: T.Tensor((16, 16), "float32", layout=T.TileLayout(T.S[(16, 16) : [8,
1]] + (T.int32(4)))),
- B: T.Tensor((16, 16), "float32", layout=T.TileLayout(T.S[(16, 16) : [8,
1]] + (T.int32(4)))),
-) -> None:
- for i, j in T.grid(16, 16):
- B[i, j] = A[i, j]
-
-
-mem_copy_m_n_p_n_m = T.dynamic("mem_copy_m_n_p_n_m", "int32")
-
-mem_copy_m_n_p_n_n = T.dynamic("mem_copy_m_n_p_n_n", "int32")
-
-
[email protected]
-def mem_copy_m_n_p_n(
- A: T.Tensor(
- (mem_copy_m_n_p_n_m, mem_copy_m_n_p_n_n),
- "float32",
- layout=T.TileLayout(
- T.S[(mem_copy_m_n_p_n_m, mem_copy_m_n_p_n_n) : [p, 1]] +
(mem_copy_m_n_p_n_n) # noqa: F821
- ),
- ),
- B: T.Tensor(
- (mem_copy_m_n_p_n_m, mem_copy_m_n_p_n_n),
- "float32",
- layout=T.TileLayout(
- T.S[(mem_copy_m_n_p_n_m, mem_copy_m_n_p_n_n) : [p, 1]] +
(mem_copy_m_n_p_n_n) # noqa: F821
- ),
- ),
- m: mem_copy_m_n_p_n_m,
- n: mem_copy_m_n_p_n_n,
- p: T.int32,
-) -> None:
- for i, j in T.grid(m, n):
- B[i, j] = A[i, j]
-
-
-def test_specialize_nothing():
- func = matmul.specialize({})
- assert func.same_as(matmul) # Pointer the same
-
-
-def test_specialize_matmul():
- a, _, _, n = matmul.params
- # fully specialized
- func = matmul.specialize(
- {a: tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((128, 128),
"float32"))}
- )
- assert_structural_equal_ignore_global_symbol(func, matmul_128)
- # partially specialized
- func = matmul.specialize({n: 128})
- assert_structural_equal_ignore_global_symbol(func, matmul_m_128)
- # symbolic specialized
- func = matmul.specialize({n: tvm_tirx.tirx.Var("x", "int32") * 8})
- assert_structural_equal_ignore_global_symbol(func, matmul_m_8x)
-
-
-def test_specialize_elemwise():
- a, c = element_wise.params
- C = c
- # fully specialized
- func = element_wise.specialize(
- {a: tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((128, 64),
"float32"))}
- )
- assert_structural_equal_ignore_global_symbol(func, element_wise_128_64)
- # partially specialized
- func = element_wise.specialize(
- {c: tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType((128,
C.ty.shape[1]), "float32"))}
- )
- assert_structural_equal_ignore_global_symbol(func, element_wise_128_n)
-
-
-def test_specialize_mem_copy():
- a, _, m, n, p, q = mem_copy.params
- # fully specialized
- func = mem_copy.specialize(
- {
- a: tvm_tirx.tirx.Var(
- "tensor",
- tvm_tirx.tirx.TensorType(
- (16, 16),
- "float32",
- layout=tvm_tirx.tirx.layout.TileLayout(
- tvm_tirx.tirx.layout.S[(16, 16) : [8, 1]] +
(T.int32(4))
- ),
- ),
- )
- }
- )
- assert_structural_equal_ignore_global_symbol(func, mem_copy_16_16_8_4)
- func = mem_copy.specialize({n: 16, m: 16, p: 8, q: 4})
- assert_structural_equal_ignore_global_symbol(func, mem_copy_16_16_8_4)
- # partially specialized
- func = mem_copy.specialize({q: n})
- assert_structural_equal_ignore_global_symbol(func, mem_copy_m_n_p_n)
-
-
-def test_specialize_recursive_load():
- # TODO(Siyuan): add recursive Load testcase, e.g. A[C[i]]
- pass
-
-
-def test_specialize_with_const_folding():
- n = T.dynamic("n", "int32")
-
- @T.function
- def before(A: T.Tensor([n // 8, 8], "int32"), B: T.Tensor([n], "int32")):
- for i in range(n - 1):
- B[i] = A[i // 8, i % 8] + (n + 1) * 42
-
- @T.function
- def expected(A: T.Tensor([2, 8], "int32"), B: T.Tensor([16], "int32")):
- for i in range(15):
- B[i] = A[i // 8, i % 8] + 714
-
- b = before.params[1]
- after = before.specialize(
- {b: tvm_tirx.tirx.Var("tensor", tvm_tirx.tirx.TensorType([16],
"int32"))}
- )
- assert_structural_equal_ignore_global_symbol(expected, after)
-
-
-def test_specialize_decl_tensor():
- """Tensors occurring in a DeclTensor statement should be updated"""
-
- @T.function(private=True)
- def before(A_data: T.handle("float32"), A_size: T.int32):
- A_buf = T.decl_tensor(A_data, ty_args=[T.Tensor(A_size, "float32",
data=A_data)])
- for i in range(A_size):
- A_buf[i] = A_buf[i] * 2.0
-
- @T.function(private=True)
- def expected(A_data: T.handle("float32")):
- A_buf = T.decl_tensor(A_data, ty_args=[T.Tensor(16, "float32",
data=A_data)])
- for i in range(16):
- A_buf[i] = A_buf[i] * 2.0
-
- param_map = {before.params[1]: T.int32(16)}
- after = before.specialize(param_map)
-
- tvm_tirx.ir.assert_structural_equal(expected, after)
-
-
-def test_specialize_preserves_decl_tensor_alias():
- before_n = T.dynamic("before_n", "int32")
-
- @T.function(private=True)
- def before(A: T.Tensor((before_n,), "int32"), n: before_n):
- A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor((n,), "int32",
data=A.data)])
- A_flat[n - 1] = 42
-
- @T.function(private=True)
- def expected(A: T.Tensor((8,), "int32")):
- A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor((8,), "int32",
data=A.data)])
- A_flat[7] = 42
-
- after = before.specialize({before.params[1]: 8})
-
- tvm_tirx.ir.assert_structural_equal(expected, after)
-
-
-def test_specialize_tensor_var_to_var():
- """A tensor var may be remapped by specialization
-
- If a tensor variable is replaced by a specialization, then other
- tensors using the same tensor var should also be updated.
- """
-
- @T.function(private=True)
- def before(A: T.Tensor([16, 16], "float32"), B: T.Tensor([16, 16],
"float32")):
- A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor([256], "float32",
data=A.data)])
- B_flat = T.decl_tensor(B.data, ty_args=[T.Tensor([256], "float32",
data=B.data)])
- for i in range(256):
- B_flat[i] = A_flat[i] * 2.0
-
- @T.function(private=True)
- def expected(A: T.Tensor([16, 16], "float32")):
- A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor([256], "float32",
data=A.data)])
- B_flat = T.decl_tensor(A.data, ty_args=[T.Tensor([256], "float32",
data=A.data)])
- for i in range(256):
- B_flat[i] = A_flat[i] * 2.0
-
- A = before.params[0]
- B_handle = before.params[1]
- param_map = {B_handle: A}
- after = before.specialize(param_map)
-
- tvm_tirx.ir.assert_structural_equal(expected, after)
-
-
-def test_specialize_tensor_var_to_expr():
- """A DeclTensor source expression may be specialized directly."""
-
- @T.function(private=True)
- def before(A_data: T.handle("float32"), B_data: T.handle("float32")):
- A_buf = T.decl_tensor(A_data, ty_args=[T.Tensor(32, "float32",
data=A_data)])
- B_buf = T.decl_tensor(B_data, ty_args=[T.Tensor(16, "float32",
data=B_data)])
- for i in range(16):
- B_buf[i] = A_buf[i] * 2.0
-
- @T.function(private=True)
- def expected(A_data: T.handle("float32")):
- A_buf = T.decl_tensor(A_data, ty_args=[T.Tensor(32, "float32",
data=A_data)])
- B_buf = T.decl_tensor(
- T.address_of(A_buf[16]), ty_args=[T.Tensor(16, "float32",
data=T.address_of(A_buf[16]))]
- )
- for i in range(16):
- B_buf[i] = A_buf[i] * 2.0
-
- B_data = before.params[1]
- # body is a SeqStmt; the first statement is DeclTensor for A_buf
- A_buf = before.body[0].var
- param_map = {B_data: tvm_tirx.tirx.address_of(A_buf[16])}
- after = before.specialize(param_map)
-
- tvm_tirx.ir.assert_structural_equal(expected, after)
-
-
-def test_specialization_updates_ty():
- """Update type in specialization
-
- A Function's native function type must reflect its specialized parameters.
- """
-
- @T.function(private=True)
- def before(n: T.int32) -> T.int32:
- return n * 10
-
- @T.function(private=True)
- def expected() -> T.int32:
- return 50
-
- ty_before = tvm_tirx.ir.FuncType([tvm_tirx.ir.PrimType("int32")],
tvm_tirx.ir.PrimType("int32"))
- tvm_tirx.ir.assert_structural_equal(before.ty, ty_before)
-
- ty_expected = tvm_tirx.ir.FuncType([], tvm_tirx.ir.PrimType("int32"))
- tvm_tirx.ir.assert_structural_equal(expected.ty, ty_expected)
-
- n = before.params[0]
- param_map = {n: 5}
- after = before.specialize(param_map)
-
- tvm_tirx.ir.assert_structural_equal(after, expected)
- tvm_tirx.ir.assert_structural_equal(after.ty, ty_expected)
-
-
-if __name__ == "__main__":
- tvm_tirx.testing.main()
diff --git a/tirx/tests/python/tirx/ir/__init__.py
b/tirx/tests/python/tirx/ir/__init__.py
new file mode 100644
index 0000000000..13a83393a9
--- /dev/null
+++ b/tirx/tests/python/tirx/ir/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git a/tirx/tests/python/tirx/ir/test_function.py
b/tirx/tests/python/tirx/ir/test_function.py
new file mode 100644
index 0000000000..5ddc966332
--- /dev/null
+++ b/tirx/tests/python/tirx/ir/test_function.py
@@ -0,0 +1,65 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test TIRx function structural mutation."""
+
+import tvm_ffi
+
+
+def test_structural_map_specializes_function_signature():
+ from tvm_tirx.ir import Evaluate, FuncType, IntImm, Return, Var,
assert_structural_equal
+ from tvm_tirx.tirx import Function, TensorType
+
+ n = Var("n", "int32")
+ scale = Var("scale", "int32")
+ tensor = Var("tensor", TensorType((n,), "float32"))
+ original = Function([tensor, scale], [Evaluate(scale), Return(tensor)],
tensor.ty)
+ bindings = {n: IntImm("int32", 16), scale: IntImm("int32", 2)}
+
+ # Removing runtime parameters is explicit. structural_map handles the shape
+ # definitions embedded in tensor types and remaps the retained tensor's
uses.
+ signature = Function([tensor], original.body, original.ret_type)
+ specialized = tvm_ffi.structural_map(
+ signature, (Var, lambda var: bindings.get(var, var)), order="pre"
+ )
+ expected_tensor = Var("tensor", TensorType((16,), "float32"))
+ expected = Function(
+ [expected_tensor],
+ [Evaluate(IntImm("int32", 2)), Return(expected_tensor)],
+ expected_tensor.ty,
+ )
+ assert_structural_equal(specialized, expected)
+ assert specialized.body.seq[1].value.same_as(specialized.params[0])
+ assert_structural_equal(
+ specialized.ty, FuncType([p.ty for p in specialized.params],
specialized.ret_type)
+ )
+ assert original.params[0].ty.shape[0].same_as(n)
+ assert original.ty.arg_types[0].shape[0].same_as(n)
+
+ # A unique declaration allows the parameter array and its Var to remain in
+ # place while a nested TensorType changes. The cached FuncType must follow.
+ def declaration():
+ return Function(
+ [Var("tensor", TensorType((n,), "float32"))], None,
TensorType((n,), "float32")
+ )
+
+ inplace = declaration()
+ inplace = tvm_ffi.structural_map(
+ inplace._move(), (Var, lambda var: bindings.get(var, var)), order="pre"
+ )
+ assert int(inplace.params[0].ty.shape[0]) == 16
+ assert int(inplace.ret_type.shape[0]) == 16
+ assert_structural_equal(inplace.ty, FuncType([inplace.params[0].ty],
inplace.ret_type))
