This is an automated email from the ASF dual-hosted git repository.
MasterJH5574 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 eafcba1c44 [Relax][TensorRT] Fix YOLO BYOC offload and partitioning
gaps (#19998)
eafcba1c44 is described below
commit eafcba1c44fffbb881bb5f838ad72e0cfafa9ce5
Author: Shushi Hong <[email protected]>
AuthorDate: Thu Jul 16 15:57:40 2026 -0400
[Relax][TensorRT] Fix YOLO BYOC offload and partitioning gaps (#19998)
Fixes #19887.
This PR fixes several Relax TensorRT BYOC issues exposed by YOLO-style
models:
- adds TensorRT support for SiLU and resize2d
- preserves operand and TupleGetItem ordering during codegen
- fixes cyclic and unsafe Tuple/TGI region merging
- handles static Shape bindings and nested packed-function outputs
- normalizes PrimType dtype arguments passed to relax.arange
With these changes, yolo11n-seg can be merged into a single TensorRT
region, while yolo11n can be imported and partitioned successfully.
---
python/tvm/relax/backend/contrib/tensorrt.py | 84 +++++++++-
python/tvm/relax/op/create.py | 6 +-
src/relax/backend/contrib/tensorrt/codegen.cc | 66 +++++---
src/relax/transform/call_tir_rewrite.cc | 138 +++++++++------
src/relax/transform/fuse_ops.cc | 142 +++++++++-------
src/relax/transform/merge_composite_functions.cc | 156 ++++++++++++++++-
src/runtime/extra/contrib/tensorrt/tensorrt_ops.cc | 99 +++++++++++
tests/python/relax/test_backend_tensorrt.py | 117 +++++++++++++
tests/python/relax/test_codegen_tensorrt.py | 121 ++++++++++++--
tests/python/relax/test_op_create.py | 5 +
tests/python/relax/test_transform.py | 44 +++++
.../relax/test_transform_fuse_ops_by_pattern.py | 76 +++++++++
.../test_transform_merge_composite_functions.py | 186 +++++++++++++++++++++
13 files changed, 1072 insertions(+), 168 deletions(-)
diff --git a/python/tvm/relax/backend/contrib/tensorrt.py
b/python/tvm/relax/backend/contrib/tensorrt.py
index 303ebc394c..21ab01b24d 100644
--- a/python/tvm/relax/backend/contrib/tensorrt.py
+++ b/python/tvm/relax/backend/contrib/tensorrt.py
@@ -23,15 +23,17 @@ converter registered under the same name (the converters
are keyed by
out of the module and annotates them for the ``tensorrt`` codegen.
"""
-from collections.abc import Mapping
-
+from tvm import tirx
from tvm.ir import IRModule
-from tvm.relax.dpl.pattern import DFPattern, is_op, wildcard
-from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions
-
-from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from tvm.relax.dpl.pattern import is_op, wildcard
+from tvm.relax.expr import ShapeExpr
+from tvm.relax.transform import (
+ FuseOpsByPattern,
+ MergeCompositeFunctions,
+ PatternCheckContext,
+)
-Pattern = tuple[str, DFPattern, Mapping[str, DFPattern]]
+from ..pattern_registry import Pattern, get_patterns_with_prefix,
register_patterns
def _op_pattern(composite_name: str, op_name: str, num_args: int) -> Pattern:
@@ -40,6 +42,70 @@ def _op_pattern(composite_name: str, op_name: str, num_args:
int) -> Pattern:
return (composite_name, is_op(op_name)(*args), {})
+def _resize2d_pattern() -> Pattern:
+ """Match the subset of resize2d that the TensorRT converter implements
exactly."""
+
+ data = wildcard()
+ size = wildcard()
+ root = is_op("relax.image.resize2d")(data, size)
+
+ def check(context: PatternCheckContext) -> bool:
+ data_expr = context.annotated_expr["data"]
+ size_expr = context.annotated_expr["size"]
+ resize = context.annotated_expr["root"]
+
+ # Follow an ANF binding when the static ShapeExpr was assigned to a
variable.
+ while size_expr in context.matched_bindings:
+ size_expr = context.matched_bindings[size_expr]
+
+ if not isinstance(size_expr, ShapeExpr) or len(size_expr.values) != 2:
+ return False
+ if not all(isinstance(dim, tirx.IntImm) and dim.value > 0 for dim in
size_expr.values):
+ return False
+
+ attrs = resize.attrs
+ if attrs.layout != "NCHW":
+ return False
+ if data_expr.ty.ndim != 4 or data_expr.ty.shape is None:
+ return False
+
+ # The runtime's DLPack-to-TensorRT dtype mapper supports only
FP16/FP32. TensorRT Resize
+ # itself preserves its input dtype, so a distinct Relax out_dtype is
not supported here.
+ if data_expr.ty.dtype not in ("float16", "float32"):
+ return False
+ if resize.ty.dtype != data_expr.ty.dtype:
+ return False
+
+ if attrs.method not in ("nearest_neighbor", "linear", "cubic"):
+ return False
+ if attrs.coordinate_transformation_mode not in (
+ "asymmetric",
+ "align_corners",
+ "half_pixel",
+ "pytorch_half_pixel",
+ ):
+ return False
+
+ # Relax `round` uses ties-to-even, for which TensorRT has no
ResizeRoundMode. The other
+ # modes map exactly to TensorRT's FLOOR, CEIL, HALF_UP, and HALF_DOWN
modes.
+ if attrs.method == "nearest_neighbor" and attrs.rounding_method not in
(
+ "floor",
+ "ceil",
+ "round_prefer_ceil",
+ "round_prefer_floor",
+ ):
+ return False
+
+ return True
+
+ return (
+ "tensorrt.image.resize2d",
+ root,
+ {"data": data, "size": size, "root": root},
+ check,
+ )
+
+
def _tensorrt_patterns() -> list[Pattern]:
patterns: list[Pattern] = []
@@ -47,6 +113,7 @@ def _tensorrt_patterns() -> list[Pattern]:
for composite, op in [
("tensorrt.nn.relu", "relax.nn.relu"),
("tensorrt.sigmoid", "relax.sigmoid"),
+ ("tensorrt.nn.silu", "relax.nn.silu"),
("tensorrt.tanh", "relax.tanh"),
("tensorrt.exp", "relax.exp"),
("tensorrt.log", "relax.log"),
@@ -92,6 +159,9 @@ def _tensorrt_patterns() -> list[Pattern]:
]:
patterns.append(_op_pattern(composite, op, 2))
+ # image.resize2d (data + target-size shape argument).
+ patterns.append(_resize2d_pattern())
+
# Convolutions and matmul (data + weight).
for composite, op in [
("tensorrt.nn.conv1d", "relax.nn.conv1d"),
diff --git a/python/tvm/relax/op/create.py b/python/tvm/relax/op/create.py
index e8c1300f01..5188140222 100644
--- a/python/tvm/relax/op/create.py
+++ b/python/tvm/relax/op/create.py
@@ -245,7 +245,7 @@ def arange(
start: PrimExprLike,
end: PrimExprLike | None = None,
step: PrimExprLike = 1,
- dtype: str | DataType | None = None,
+ dtype: str | DataType | PrimType | None = None,
) -> Expr:
"""Construct a tensor with evenly spaced elements.
@@ -261,7 +261,7 @@ def arange(
step : PrimExprLike
The step size.
- dtype : Optional[str | DataType]
+ dtype : Optional[str | DataType | PrimType]
The data type of the created tensor.
Returns
@@ -288,7 +288,7 @@ def arange(
start = prim_value(start)
end = prim_value(end)
step = prim_value(step)
- return _ffi_api.arange(start, end, step, dtype) # type: ignore
+ return _ffi_api.arange(start, end, step, _raw_dtype(dtype)) # type: ignore
def hamming_window(window_size, periodic, alpha, beta, dtype):
diff --git a/src/relax/backend/contrib/tensorrt/codegen.cc
b/src/relax/backend/contrib/tensorrt/codegen.cc
index b09c988931..2b5d4d3db7 100644
--- a/src/relax/backend/contrib/tensorrt/codegen.cc
+++ b/src/relax/backend/contrib/tensorrt/codegen.cc
@@ -26,11 +26,13 @@
#include <tvm/ir/module.h>
#include <tvm/ir/op.h>
#include <tvm/ir/transform.h>
+#include <tvm/relax/analysis.h>
#include <tvm/relax/attrs/manipulate.h>
#include <tvm/relax/attrs/nn.h>
#include <tvm/relax/attrs/statistical.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
+#include <tvm/relax/utils.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/index_map.h>
@@ -97,18 +99,13 @@ using JSONGraphObjectPtr =
backend::contrib::JSONGraphObjectPtr;
using OpAttrExtractor = backend::contrib::OpAttrExtractor;
using JSONSerializer = backend::contrib::JSONSerializer;
-class TensorRTJSONSerializer;
-
/*!
- * \brief Collect the constants and attributes from all operator calls in the
body
- * of a "Composite" function.
+ * \brief Collect the primitive operator call and its attributes from a
"Composite" function.
*/
class CollectFromCompositeFunctionBody : public ExprVisitor {
public:
- explicit CollectFromCompositeFunctionBody(TensorRTJSONSerializer* serializer)
- : serializer_(serializer), node_(std::make_shared<JSONGraphNode>()) {}
+ CollectFromCompositeFunctionBody() :
node_(std::make_shared<JSONGraphNode>()) {}
- void VisitExpr_(const ConstantNode* constant_node) final;
void VisitExpr_(const CallNode* call_node) final;
void SetGenericAttributes(const CallNode* call_node) {
@@ -212,9 +209,8 @@ class CollectFromCompositeFunctionBody : public ExprVisitor
{
return true;
}
- TensorRTJSONSerializer* serializer_;
- /*! \brief Accumulated translated arguments. */
- std::vector<JSONGraphNodeEntry> args_;
+ /*! \brief The primitive operator call in the composite function. */
+ const CallNode* operator_call_{nullptr};
/*!
* \brief Temporary node into which we'll accumulate attributes. Ideally
this would be the
* final JSONGraphNode however we don't yet know how many inputs that will
have.
@@ -245,20 +241,37 @@ class TensorRTJSONSerializer : public JSONSerializer {
TVM_FFI_ICHECK(opt_composite.has_value());
std::string name = opt_composite.value();
- // Collect the constants and attributes of all operator calls inside the
composite body.
- CollectFromCompositeFunctionBody collector(this);
+ // TensorRT patterns describe a single primitive operator and use the
Composite function only
+ // to bind that operator's arguments and attributes.
+ CollectFromCompositeFunctionBody collector;
collector.VisitExpr(fn->body);
-
- // Capture the args to the "Composite" function as inputs for this node.
- std::vector<JSONGraphNodeEntry> inputs;
- for (const auto& arg : call_node->args) {
- auto res = VisitExpr(arg);
- inputs.insert(inputs.end(), res.begin(), res.end());
+ TVM_FFI_ICHECK(collector.operator_call_ != nullptr)
+ << "TensorRT Composite function " << name
+ << " must contain exactly one primitive Relax operator call";
+
+ // Bind Composite parameters back to the caller. The primitive operator's
tensor arguments
+ // are serialized below in their original order, regardless of whether
they are constants or
+ // parameters. Non-tensor scalar and shape arguments have already been
captured as attributes.
+ TVM_FFI_ICHECK_EQ(fn->params.size(), call_node->args.size());
+ ffi::Map<Var, Expr> param_bindings;
+ for (size_t i = 0; i < call_node->args.size(); ++i) {
+ param_bindings.Set(fn->params[i], call_node->args[i]);
}
- // Capture constants from the composite function body as additional inputs
for this node.
- for (const auto& node : collector.args_) {
- inputs.emplace_back(node);
+ std::vector<JSONGraphNodeEntry> inputs;
+ auto append_tensor_inputs = [&](const auto& self, const Expr& expr) ->
void {
+ if (const auto* tuple = expr.as<TupleNode>()) {
+ for (const Expr& field : tuple->fields) self(self, field);
+ return;
+ }
+ Type type = GetType(expr);
+ if (type->IsInstance<TensorTypeNode>() ||
type->IsInstance<TupleTypeNode>()) {
+ auto entries = VisitExpr(expr);
+ inputs.insert(inputs.end(), entries.begin(), entries.end());
+ }
+ };
+ for (const Expr& arg : collector.operator_call_->args) {
+ append_tensor_inputs(append_tensor_inputs, Bind(arg, param_bindings));
}
// Create the final node.
@@ -299,13 +312,12 @@ class TensorRTJSONSerializer : public JSONSerializer {
ffi::Map<Var, Expr> bindings_;
};
-void CollectFromCompositeFunctionBody::VisitExpr_(const ConstantNode*
constant_node) {
- for (const auto& entry :
serializer_->VisitExpr(ffi::GetRef<Constant>(constant_node))) {
- args_.emplace_back(entry);
- }
-}
-
void CollectFromCompositeFunctionBody::VisitExpr_(const CallNode* call_node) {
+ TVM_FFI_ICHECK(call_node->op->IsInstance<OpNode>())
+ << "TensorRT Composite functions must contain exactly one primitive
Relax operator call";
+ TVM_FFI_ICHECK(operator_call_ == nullptr)
+ << "TensorRT Composite functions must contain exactly one primitive
Relax operator call";
+ operator_call_ = call_node;
if (!TrySetLayoutTransformAttributes(call_node)) {
SetGenericAttributes(call_node);
SetArgumentAttributes(call_node);
diff --git a/src/relax/transform/call_tir_rewrite.cc
b/src/relax/transform/call_tir_rewrite.cc
index e1d9a265dc..20c7027c97 100644
--- a/src/relax/transform/call_tir_rewrite.cc
+++ b/src/relax/transform/call_tir_rewrite.cc
@@ -29,6 +29,8 @@
#include <tvm/relax/type.h>
#include <tvm/tirx/op.h>
+#include <algorithm>
+
#include "utils.h"
namespace tvm {
@@ -73,27 +75,12 @@ class CallTIRMutator : public ExprMutator {
bool is_inplace = call->op.same_as(call_tir_inplace_op);
const auto* inplace_attrs = call->attrs.as<CallTIRInplaceAttrs>();
ffi::Array<Expr> outs;
+ ffi::Optional<TupleType> tuple_output_type;
if (const auto& tensor_ty = MatchType<TensorType>(expr)) {
// single output case
const TensorType& output_ty = tensor_ty.value();
- TVM_FFI_ICHECK(output_ty->shape.has_value())
- << "the TensorType shape of call_tir has not populated";
- int dev_index = 0;
- ffi::String scope = "global";
- if (output_ty->vdevice.has_value()) {
- dev_index = GetDeviceIndex(mod_, output_ty->vdevice.value());
- scope = output_ty->vdevice.value()->memory_scope;
- } else {
- dev_index = GetDeviceIndexByScope(mod_, scope);
- }
-
if (!is_inplace) {
- outs.push_back(builder_->Emit(Call(Type::Missing(), alloc_tensor_op,
-
{output_ty->shape.value().as_or_throw<ShapeExpr>(),
-
DataTypeImm(output_ty->dtype.value()->dtype),
- IntImm::Int64(dev_index),
StringImm(scope)},
- Attrs(), {output_ty}),
- "alloc"));
+ outs.push_back(AllocateOutputTensor(output_ty, alloc_tensor_op));
} else {
// if there is only one output, it must be an in-place argument, but
check anyway
TVM_FFI_ICHECK(inplace_attrs->inplace_indices[0] != -1)
@@ -105,35 +92,30 @@ class CallTIRMutator : public ExprMutator {
} else if (const auto& tuple_ty = MatchType<TupleType>(expr)) {
// multiple output case
const TupleType& output_ty = tuple_ty.value();
- for (size_t i = 0; i < output_ty->fields.size(); ++i) {
- const auto& field = output_ty->fields[i];
-
- TVM_FFI_ICHECK(field->IsInstance<TensorTypeNode>())
- << "call_tir expects Tuple of TensorType, but got " << field
- << " as an element of TupleType";
- const auto& field_tensor = field.as_or_throw<TensorType>();
- TVM_FFI_ICHECK(field_tensor->shape.has_value())
- << "call_tir expects all TensorType has shape, but got " <<
field_tensor
- << " as an element of TupleType";
-
- int dev_index = 0;
- ffi::String scope = "global";
- if (field_tensor->vdevice.has_value()) {
- dev_index = GetDeviceIndex(mod_, field_tensor->vdevice.value());
- scope = field_tensor->vdevice.value()->memory_scope;
- }
-
- if (!is_inplace || inplace_attrs->inplace_indices[i] == -1) {
- outs.push_back(
- builder_->Emit(Call(Type::Missing(), alloc_tensor_op,
-
{field_tensor->shape.value().as_or_throw<ShapeExpr>(),
-
DataTypeImm(field_tensor->dtype.value()->dtype),
- IntImm::Int64(dev_index),
StringImm(scope)},
- Attrs(), {field_tensor}),
- "alloc"));
- } else {
- outs.push_back(
-
call->args[1].as_or_throw<Tuple>()->fields[inplace_attrs->inplace_indices[i]]);
+ tuple_output_type = output_ty;
+ bool has_nested_tuple =
+ std::any_of(output_ty->fields.begin(), output_ty->fields.end(),
+ [](const Type& field) { return
field->IsInstance<TupleTypeNode>(); });
+
+ if (has_nested_tuple) {
+ TVM_FFI_ICHECK(!is_inplace)
+ << "call_tir_inplace does not support nested tuple output types";
+ FlattenAndAllocateOutputs(output_ty, alloc_tensor_op, &outs);
+ } else {
+ for (size_t i = 0; i < output_ty->fields.size(); ++i) {
+ const auto& field = output_ty->fields[i];
+
+ TVM_FFI_ICHECK(field->IsInstance<TensorTypeNode>())
+ << "call_tir expects Tuple of TensorType, but got " << field
+ << " as an element of TupleType";
+ const auto& field_tensor = field.as_or_throw<TensorType>();
+
+ if (!is_inplace || inplace_attrs->inplace_indices[i] == -1) {
+ outs.push_back(AllocateOutputTensor(field_tensor,
alloc_tensor_op));
+ } else {
+ outs.push_back(
+
call->args[1].as_or_throw<Tuple>()->fields[inplace_attrs->inplace_indices[i]]);
+ }
}
}
} else {
@@ -166,15 +148,75 @@ class CallTIRMutator : public ExprMutator {
builder_->Emit(Call(Type::Missing(), call->args[0], args), "_");
}
- if (outs.size() == 1) {
- return outs[0];
+ if (tuple_output_type.has_value()) {
+ size_t index = 0;
+ Expr output = RebuildOutputTuple(tuple_output_type.value(), outs,
&index);
+ TVM_FFI_ICHECK_EQ(index, outs.size());
+ return output;
}
+ if (outs.size() == 1) return outs[0];
return std::move(Tuple(outs));
}
return ffi::GetRef<Expr>(call);
}
+ Expr AllocateOutputTensor(const TensorType& tensor_ty, const Op&
alloc_tensor_op) {
+ TVM_FFI_ICHECK(tensor_ty->shape.has_value())
+ << "call_tir expects all TensorType has shape, but got " << tensor_ty;
+
+ int dev_index = 0;
+ ffi::String scope = "global";
+ if (tensor_ty->vdevice.has_value()) {
+ dev_index = GetDeviceIndex(mod_, tensor_ty->vdevice.value());
+ scope = tensor_ty->vdevice.value()->memory_scope;
+ } else {
+ dev_index = GetDeviceIndexByScope(mod_, scope);
+ }
+
+ return builder_->Emit(Call(Type::Missing(), alloc_tensor_op,
+
{tensor_ty->shape.value().as_or_throw<ShapeExpr>(),
+ DataTypeImm(tensor_ty->dtype.value()->dtype),
+ IntImm::Int64(dev_index), StringImm(scope)},
+ Attrs(), {tensor_ty}),
+ "alloc");
+ }
+
+ void FlattenAndAllocateOutputs(const Type& type, const Op& alloc_tensor_op,
+ ffi::Array<Expr>* outs) {
+ if (const auto* tensor_ty = type.as<TensorTypeNode>()) {
+ outs->push_back(AllocateOutputTensor(ffi::GetRef<TensorType>(tensor_ty),
alloc_tensor_op));
+ return;
+ }
+ if (const auto* tuple_ty = type.as<TupleTypeNode>()) {
+ for (const Type& field : tuple_ty->fields) {
+ FlattenAndAllocateOutputs(field, alloc_tensor_op, outs);
+ }
+ return;
+ }
+ TVM_FFI_THROW(TypeError) << "call_tir expects nested tuple outputs to
contain only "
+ "TensorType, but got "
+ << type;
+ }
+
+ Expr RebuildOutputTuple(const Type& type, const ffi::Array<Expr>& outs,
size_t* index) {
+ if (type->IsInstance<TensorTypeNode>()) {
+ TVM_FFI_ICHECK_LT(*index, outs.size());
+ return outs[(*index)++];
+ }
+ if (const auto* tuple_ty = type.as<TupleTypeNode>()) {
+ ffi::Array<Expr> fields;
+ fields.reserve(tuple_ty->fields.size());
+ for (const Type& field : tuple_ty->fields) {
+ fields.push_back(RebuildOutputTuple(field, outs, index));
+ }
+ return Tuple(std::move(fields));
+ }
+ TVM_FFI_THROW(TypeError) << "call_tir expects nested tuple outputs to
contain only "
+ "TensorType, but got "
+ << type;
+ }
+
/*! \brief The context IRModule. */
IRModule mod_;
};
diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc
index b76e129043..7d8986218b 100644
--- a/src/relax/transform/fuse_ops.cc
+++ b/src/relax/transform/fuse_ops.cc
@@ -396,7 +396,8 @@ class GraphCreator : public ExprVisitor {
*/
class FunctionCreator : public ExprMutator {
public:
- explicit FunctionCreator(bool lift_constant) : lift_constant_(lift_constant)
{}
+ explicit FunctionCreator(bool lift_constant, ffi::Map<Var, Expr>
outer_bindings)
+ : outer_bindings_(std::move(outer_bindings)),
lift_constant_(lift_constant) {}
/*!
* \brief Append a new binding to this function and possibly create new
parameters for the
* function accordingly
@@ -475,16 +476,15 @@ class FunctionCreator : public ExprMutator {
}
/*! \brief Set a var defined in the group as output. */
- size_t AppendOutput(const Var& var) {
+ void AppendOutput(const Var& var) {
TVM_FFI_ICHECK(defined_vars_.count(var.get()));
- auto output_idx = GetOutputIndex(var);
- if (output_idx) {
- return *output_idx;
- }
+ if (GetOutputIndex(var)) return;
output_vars_.push_back(var.get());
- return output_vars_.size() - 1;
}
+ /*! \brief Variables returned from the grouped function, in return-value
order. */
+ const std::vector<const VarNode*>& output_vars() const { return
output_vars_; }
+
/*!
* \brief Create the grouped function according to the collected bindings
and parameters
* \param composite_name The name to identify the pattern this function is
created from, if any.
@@ -622,6 +622,21 @@ class FunctionCreator : public ExprMutator {
// If the expression is not a variable or is a undefined variable, it
should be populated as a
// parameter of the relax function.
const auto* var = expr.as<VarNode>();
+ if (var != nullptr && defined_vars_.count(var) == 0) {
+ Var bound_var = ffi::GetRef<Var>(var);
+ Expr bound_value = bound_var;
+ std::unordered_set<const VarNode*> visited;
+ while (const auto* current_var = bound_value.as<VarNode>()) {
+ if (!visited.insert(current_var).second) break;
+ auto it = outer_bindings_.find(ffi::GetRef<Var>(current_var));
+ if (it == outer_bindings_.end()) break;
+ bound_value = (*it).second;
+ }
+ if (!bound_value.same_as(bound_var) &&
IsInlinableConstants(bound_value)) {
+ inlined_bindings_[var] = bound_value;
+ return;
+ }
+ }
if ((var == nullptr || defined_vars_.count(var) == 0) &&
(lift_constant_ || !expr->IsInstance<ConstantNode>())) {
ffi::String name = var != nullptr
@@ -650,6 +665,11 @@ class FunctionCreator : public ExprMutator {
if (it != arguments_.end()) {
return params_[it - arguments_.begin()];
}
+ if (const auto* var = expr.as<VarNode>()) {
+ if (auto inlined = inlined_bindings_.find(var); inlined !=
inlined_bindings_.end()) {
+ return inlined->second;
+ }
+ }
// Otherwise, recurse into this expression.
return ExprMutator::VisitExpr(expr);
}
@@ -674,10 +694,14 @@ class FunctionCreator : public ExprMutator {
private:
/*! \brief The variables defined in this function */
std::unordered_set<const VarNode*> defined_vars_;
+ /*! \brief Caller variables replaced by statically inlinable bound values. */
+ std::unordered_map<const VarNode*, Expr> inlined_bindings_;
/*! \brief The number of parameters reserved for constants */
int n_param_for_const_ = 0;
/*! \brief The output vars */
std::vector<const VarNode*> output_vars_;
+ /*! \brief Bindings in the caller function, used to inline static leaf
expressions. */
+ ffi::Map<Var, Expr> outer_bindings_;
/*! \brief Whether or not to lift bound constants to parameters */
bool lift_constant_;
/*! \brief Mapping from tuple parameter of the function to its position
index */
@@ -750,8 +774,10 @@ class OperatorFusor : public ExprMutator {
// attr::kCodegen.
if (func->IsInstance<relax::FunctionNode>() &&
!func->HasNonzeroAttr(attr::kPrimitive) &&
!func->GetAttr<ffi::String>(attr::kCodegen).has_value()) {
+ outer_bindings_ = AnalyzeVar2Value(func);
auto updated_func = VisitExpr(func).as_or_throw<Function>();
builder_->UpdateFunction(gv, updated_func);
+ outer_bindings_ = {};
}
}
return builder_->GetContextIRModule();
@@ -770,22 +796,6 @@ class OperatorFusor : public ExprMutator {
return obj2group;
}
- bool IsTupleOutput(Function f) {
- auto ty = GetType(f).as<FuncTypeNode>();
- TVM_FFI_ICHECK(ty);
- return ty->ret->IsInstance<TupleTypeNode>();
- }
-
- bool IsNestedTupleOutput(Function f) {
- if (!IsTupleOutput(f)) return false;
-
- auto tup = GetType(f).as<FuncTypeNode>()->ret.as<TupleTypeNode>();
- for (const auto& field : tup->fields) {
- if (field->IsInstance<TupleTypeNode>()) return true;
- }
- return false;
- }
-
BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) final {
group2func_.clear();
@@ -808,9 +818,9 @@ class OperatorFusor : public ExprMutator {
// last binding of the group.
builder_->BeginDataflowBlock();
- // For each group, record which variables need to be remapped to the
output of TupleGetItem.
- // Only relevant when the output of the grouped function is a tuple.
- std::unordered_map<Group*, std::vector<Var>> pending_tuple_get;
+ // Preserve the original binding order when emitting TupleGetItem bindings
for groups with
+ // multiple boundary outputs. Missing entries are filled when the grouped
call is emitted.
+ std::unordered_map<Group*, std::vector<Var>> pending_output_remap;
// A grouped function which returns a tuple requires attaching
TupleGetItem to each element and
// remapping variables in earlier bindings appropriately. Thus, a binding
whose value depends on
@@ -837,15 +847,10 @@ class OperatorFusor : public ExprMutator {
}
const Function& func = func_info.function_.value();
- // If this binding belongs to a group whose output is a tuple, the
original bound variable
- // needs to be remapped to the output of TupleGetItem after the
corresponding tuple is
- // emitted.
- if (IsTupleOutput(func) && tuple_get_indices_.count(binding->var.get()))
{
- if (!GetType(binding->var)->IsInstance<TupleTypeNode>() ||
IsNestedTupleOutput(func)) {
- // When binding->var itself is a tuple, we do not need to remap this
variable to the
- // output of TupleGetItem unless the output is a nested tuple.
- pending_tuple_get[group].push_back(binding->var);
- }
+ const auto& output_vars = func_info.output_vars();
+ if (output_vars.size() > 1 && std::find(output_vars.begin(),
output_vars.end(),
+ binding->var.get()) !=
output_vars.end()) {
+ pending_output_remap[group].push_back(binding->var);
}
// Case 2. If the binding is not the last binding of the group, we skip
it.
@@ -863,29 +868,50 @@ class OperatorFusor : public ExprMutator {
GlobalVar gv = builder_->AddFunction(func, func_info.name_hint_);
// Step b. Create the call to the deduplicated function, and then emit
the call.
- // - If this binding is an output binding, emit an output variable.
- // - Otherwise, emit a dataflow variable.
+ // A multi-output call is internal to this dataflow block, while a
single-output call has the
+ // same dataflow/output status as its sole boundary variable. The last
binding is only the
+ // insertion point and may itself be a dead internal binding.
+ TVM_FFI_ICHECK(!output_vars.empty());
Var new_var;
Call call_to_emit = Call(Type::Missing(), gv,
UpdateArgs(func_info.arguments_));
- if (var_binding->var->IsInstance<DataflowVarNode>()) {
- new_var = builder_->Emit(call_to_emit);
- } else {
+ if (output_vars.size() == 1 &&
!output_vars[0]->IsInstance<DataflowVarNode>()) {
new_var = builder_->EmitOutput(call_to_emit);
+ } else {
+ new_var = builder_->Emit(call_to_emit);
}
- // Step c. Update the mapping used for the remapping of the binding
variables.
- if (IsTupleOutput(func) && !pending_tuple_get.empty()) {
- // If the output is a tuple, attach TupleGetItem to all tuple
elements, and
- // remap variables approriately.
- // The variables that need to be remapped and the corresponding tuple
indices are
- // available in pending_tuple_get and tuple_get_indices_ respectively.
- for (const auto& var : pending_tuple_get[group]) {
- auto tuple_get = TupleGetItem(new_var,
tuple_get_indices_[var.get()]);
- var_remap_[var] = builder_->Emit(tuple_get);
+ // Step c. Remap every boundary output to the corresponding result of
the grouped call.
+ // FunctionCreator uses output_vars() order when it constructs a
multi-output tuple. A
+ // single boundary output is returned directly, including when that
output is itself a tuple.
+ if (output_vars.size() == 1) {
+ var_remap_[ffi::GetRef<Var>(output_vars[0])] = new_var;
+ continue;
+ }
+
+ std::unordered_set<const VarNode*> remapped_outputs;
+ auto remap_output = [&](const Var& output_var) {
+ auto it = std::find(output_vars.begin(), output_vars.end(),
output_var.get());
+ TVM_FFI_ICHECK(it != output_vars.end());
+ int index = static_cast<int>(std::distance(output_vars.begin(), it));
+ TupleGetItem tuple_get(new_var, index);
+ if (output_var->IsInstance<DataflowVarNode>()) {
+ var_remap_[output_var] = builder_->Emit(tuple_get);
+ } else {
+ var_remap_[output_var] = builder_->EmitOutput(tuple_get);
+ }
+ remapped_outputs.insert(output_var.get());
+ };
+
+ if (auto it = pending_output_remap.find(group); it !=
pending_output_remap.end()) {
+ for (const Var& output_var : it->second) {
+ remap_output(output_var);
+ }
+ }
+ for (const VarNode* output_var : output_vars) {
+ if (!remapped_outputs.count(output_var)) {
+ remap_output(ffi::GetRef<Var>(output_var));
}
- } else {
- var_remap_[var_binding->var] = new_var;
}
}
// Step 5. Finish the binding block generation.
@@ -907,10 +933,8 @@ class OperatorFusor : public ExprMutator {
}
// Add the binding to the grouped function it's in, and update the
function information
// accordingly.
- if (!group2func_.count(group)) {
- group2func_.emplace(group, lift_constants_);
- }
- group2func_.find(group)->second.AppendBinding(binding);
+ auto it = group2func_.try_emplace(group, lift_constants_,
outer_bindings_).first;
+ it->second.AppendBinding(binding);
}
}
@@ -940,8 +964,7 @@ class OperatorFusor : public ExprMutator {
if (auto producer = group2func_.find(producer_group);
producer_group != cur_group && producer != group2func_.end()) {
- auto output_index = producer->second.AppendOutput(used_var);
- tuple_get_indices_[used_var.get()] = output_index;
+ producer->second.AppendOutput(used_var);
}
}
};
@@ -1040,9 +1063,8 @@ class OperatorFusor : public ExprMutator {
GroupMap obj2group_;
/*! \brief Internal function information map. */
std::unordered_map<Group*, FunctionCreator> group2func_;
- /*! \brief Record the index for TupleGetItem if the variable needs to be
remapped to an output
- * tuple element after fusion. */
- std::unordered_map<const VarNode*, int> tuple_get_indices_;
+ /*! \brief Bindings visible while rewriting the current Relax function. */
+ ffi::Map<Var, Expr> outer_bindings_;
/*!
* \brief A map from a group to its dependent groups, used to detect cyclic
dependencies.
* \note Use vector so we can be deterministic, there won't be a lot of dep
groups so
diff --git a/src/relax/transform/merge_composite_functions.cc
b/src/relax/transform/merge_composite_functions.cc
index 18e2736ecb..90a8bc2b95 100644
--- a/src/relax/transform/merge_composite_functions.cc
+++ b/src/relax/transform/merge_composite_functions.cc
@@ -56,6 +56,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/relax/analysis.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
@@ -81,17 +82,19 @@ class CompositeGroupsBuilder : public
MemoizedExprTranslator<Group*> {
CompositeGroupsBuilder(IRModule mod, support::Arena* arena) : mod_(mod),
arena_(arena) {}
GroupMap Run(Function func) {
+ var_usage_ = CollectVarUsage(func);
+ for (const auto& [var, value] : var_usage_.bound_values) {
+ value_to_bound_vars_[value.get()].push_back(var);
+ }
+
for (const auto& param : func->params) {
memo_[param] = arena_->make<Group>();
}
- PostOrderVisit(func, [this](Expr e) {
- // Make default groups for dataflow nodes other than CallNode.
- // Groups for CallNode are created in its visitor.
- if (e->IsInstance<ConstantNode>() || e->IsInstance<ShapeExprNode>() ||
- e->IsInstance<TupleNode>() || e->IsInstance<TupleGetItemNode>() ||
- (!e->IsInstance<CallNode>() && !e->IsInstance<VarNode>() &&
e.as<PrimExpr>())) {
- memo_[e] = arena_->make<Group>();
+ PostOrderVisit(func, [this](const Expr& expr) {
+ if (expr->IsInstance<ConstantNode>() ||
expr->IsInstance<ShapeExprNode>() ||
+ (!expr->IsInstance<CallNode>() && !expr->IsInstance<VarNode>() &&
expr.as<PrimExpr>())) {
+ memo_[expr] = arena_->make<Group>();
}
});
@@ -143,6 +146,9 @@ class CompositeGroupsBuilder : public
MemoizedExprTranslator<Group*> {
}
Group* VisitExpr_(const CallNode* call) {
+ for (const Expr& arg : call->args) {
+ EnsureVisited(arg);
+ }
std::vector<Group*> groups_to_merge = GetGroupsToMerge(call);
Group* group;
@@ -166,7 +172,65 @@ class CompositeGroupsBuilder : public
MemoizedExprTranslator<Group*> {
return group;
}
+ Group* VisitExpr_(const TupleNode* tuple) {
+ Expr tuple_expr = ffi::GetRef<Tuple>(tuple);
+ if (!IsFlatTensorTuple(tuple_expr)) return arena_->make<Group>();
+
+ for (const Expr& field : tuple->fields) {
+ EnsureVisited(field);
+ }
+ if (HasOnlyTupleGetItemUsers(tuple_expr)) {
+ Group* tuple_group = nullptr;
+ for (const Expr& field : tuple->fields) {
+ auto it = memo_.find(field);
+ if (it == memo_.end()) {
+ tuple_group = nullptr;
+ break;
+ }
+ Group* field_group = it->second->FindRoot();
+ if (tuple_group == nullptr) {
+ tuple_group = field_group;
+ } else if (tuple_group != field_group) {
+ tuple_group = nullptr;
+ break;
+ }
+ }
+ if (tuple_group != nullptr && CanAbsorbTupleNodes(tuple_group)) {
+ tuple_group->num_nodes += 1;
+ return tuple_group;
+ }
+ }
+
+ Group* group = arena_->make<Group>();
+ UpdateGroupDependencies(group, tuple->fields);
+ return group;
+ }
+
+ Group* VisitExpr_(const TupleGetItemNode* tuple_get_item) {
+ if (!IsFlatTensorTuple(tuple_get_item->tuple)) return
arena_->make<Group>();
+
+ EnsureVisited(tuple_get_item->tuple);
+ auto it = memo_.find(tuple_get_item->tuple);
+ if (it != memo_.end() && HasOnlyTupleGetItemUsers(tuple_get_item->tuple)) {
+ Group* tuple_group = it->second->FindRoot();
+ if (CanAbsorbTupleNodes(tuple_group)) {
+ tuple_group->num_nodes += 1;
+ return tuple_group;
+ }
+ }
+
+ Group* group = arena_->make<Group>();
+ UpdateGroupDependencies(group, {tuple_get_item->tuple});
+ return group;
+ }
+
private:
+ void EnsureVisited(const Expr& expr) {
+ if (!expr.as<GlobalVarNode>() && !memo_.count(expr)) {
+ VisitExpr(expr);
+ }
+ }
+
ffi::Optional<ffi::String> GetCodegenName(const Expr& callee) {
auto const* gvar = callee.as<GlobalVarNode>();
if (!gvar) {
@@ -189,6 +253,79 @@ class CompositeGroupsBuilder : public
MemoizedExprTranslator<Group*> {
return std::nullopt;
}
+ bool CanAbsorbTupleNodes(Group* group) { return
GetCodegenName(group->FindRoot()).has_value(); }
+
+ bool IsFlatTensorTuple(const Expr& expr) {
+ const auto* tuple_type = GetType(expr).as<TupleTypeNode>();
+ if (tuple_type == nullptr || tuple_type->fields.empty()) return false;
+ return std::all_of(tuple_type->fields.begin(), tuple_type->fields.end(),
+ [](const Type& field) { return
field->IsInstance<TensorTypeNode>(); });
+ }
+
+ bool HasOnlyTupleGetItemUsers(const Expr& tuple_expr) {
+ if (auto it = tuple_get_item_only_usage_.find(tuple_expr.get());
+ it != tuple_get_item_only_usage_.end()) {
+ return it->second;
+ }
+
+ bool result = ComputeHasOnlyTupleGetItemUsers(tuple_expr);
+ tuple_get_item_only_usage_[tuple_expr.get()] = result;
+ return result;
+ }
+
+ bool ComputeHasOnlyTupleGetItemUsers(const Expr& tuple_expr) {
+ ffi::Optional<Var> tuple_var;
+ if (const auto* var = tuple_expr.as<VarNode>()) {
+ tuple_var = ffi::GetRef<Var>(var);
+ } else if (auto it = value_to_bound_vars_.find(tuple_expr.get());
+ it != value_to_bound_vars_.end() && it->second.size() == 1) {
+ tuple_var = it->second[0];
+ }
+ if (!tuple_var.has_value()) return false;
+
+ // Follow aliases in both directions. Checking only tuple_var would allow
an alias to be used
+ // exclusively by TupleGetItem while the original tuple still escapes from
the external region.
+ std::vector<Var> pending{tuple_var.value()};
+ std::unordered_set<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> aliases;
+ bool has_tuple_get_item = false;
+ while (!pending.empty()) {
+ Var current = pending.back();
+ pending.pop_back();
+ if (!aliases.insert(current).second) continue;
+
+ if (std::any_of(var_usage_.outputs.begin(), var_usage_.outputs.end(),
+ [&](const Var& output) { return output.same_as(current);
})) {
+ return false;
+ }
+
+ // Walk toward the original tuple value when current is itself an alias.
+ if (auto binding_it = var_usage_.bound_values.find(current);
+ binding_it != var_usage_.bound_values.end()) {
+ if (const auto* source = (*binding_it).second.as<VarNode>()) {
+ pending.push_back(ffi::GetRef<Var>(source));
+ }
+ }
+
+ auto uses_it = var_usage_.downstream_usage.find(current);
+ if (uses_it == var_usage_.downstream_usage.end()) continue;
+ for (const Var& user : (*uses_it).second) {
+ auto binding_it = var_usage_.bound_values.find(user);
+ if (binding_it == var_usage_.bound_values.end()) return false;
+ const Expr& bound_value = (*binding_it).second;
+ if (const auto* tuple_get_item = bound_value.as<TupleGetItemNode>()) {
+ if (!tuple_get_item->tuple.same_as(current)) return false;
+ has_tuple_get_item = true;
+ } else if (const auto* source = bound_value.as<VarNode>()) {
+ if (!ffi::GetRef<Var>(source).same_as(current)) return false;
+ pending.push_back(user);
+ } else {
+ return false;
+ }
+ }
+ }
+ return has_tuple_get_item;
+ }
+
Group* CreateNewGroup(const CallNode* call) {
Group* group = arena_->make<Group>();
if (ffi::Optional<ffi::String> codegen_name = GetCodegenName(call->op)) {
@@ -226,6 +363,7 @@ class CompositeGroupsBuilder : public
MemoizedExprTranslator<Group*> {
std::unordered_set<Group*> dependencies;
for (const auto& arg : args) {
+ if (arg.as<GlobalVarNode>()) continue;
for (auto dep : group_deps_[memo_[arg]->FindRoot()]) {
dependencies.insert(dep);
}
@@ -279,6 +417,7 @@ class CompositeGroupsBuilder : public
MemoizedExprTranslator<Group*> {
std::unordered_set<Group*> parent_dependencies =
GetParentGroupDependencies(call->args);
for (const auto& arg : call->args) {
+ if (arg.as<GlobalVarNode>()) continue;
auto arg_group = memo_[arg];
ffi::Optional<ffi::String> arg_codegen_name = GetCodegenName(arg_group);
if (arg_codegen_name == codegen_name &&
!parent_dependencies.count(arg_group->FindRoot())) {
@@ -294,6 +433,9 @@ class CompositeGroupsBuilder : public
MemoizedExprTranslator<Group*> {
IRModule mod_;
support::Arena* arena_;
+ VarUsageInfo var_usage_;
+ std::unordered_map<const ffi::Object*, std::vector<Var>>
value_to_bound_vars_;
+ std::unordered_map<const ffi::Object*, bool> tuple_get_item_only_usage_;
// Map from group to its dependencies. All groups in this map, whether it's
// the key or in value, should be root node (that is, group->parent ==
nullptr).
std::unordered_map<Group*, std::unordered_set<Group*>> group_deps_;
diff --git a/src/runtime/extra/contrib/tensorrt/tensorrt_ops.cc
b/src/runtime/extra/contrib/tensorrt/tensorrt_ops.cc
index 61dee5a0fb..0004b19124 100644
--- a/src/runtime/extra/contrib/tensorrt/tensorrt_ops.cc
+++ b/src/runtime/extra/contrib/tensorrt/tensorrt_ops.cc
@@ -195,6 +195,23 @@ class ActivationOpConverter : public TensorRTOpConverter {
}
};
+class SiluOpConverter : public TensorRTOpConverter {
+ public:
+ explicit SiluOpConverter(std::string op_name)
+ : TensorRTOpConverter(std::move(op_name), {kTensor}) {}
+ ~SiluOpConverter() = default;
+
+ void Convert(TensorRTOpConverterParams* params) const {
+ auto input = params->inputs.at(0).tensor;
+ auto sigmoid_layer = params->network->addActivation(*input,
nvinfer1::ActivationType::kSIGMOID);
+ TVM_FFI_ICHECK(sigmoid_layer != nullptr);
+ auto silu_layer = params->network->addElementWise(*input,
*sigmoid_layer->getOutput(0),
+
nvinfer1::ElementWiseOperation::kPROD);
+ TVM_FFI_ICHECK(silu_layer != nullptr);
+ params->outputs.push_back(silu_layer->getOutput(0));
+ }
+};
+
class ElementWiseBinaryOpConverter : public TensorRTOpConverter {
public:
explicit ElementWiseBinaryOpConverter(std::string op_name)
@@ -1237,6 +1254,84 @@ class StridedSliceOpConverter : public
TensorRTOpConverter {
};
#endif
+#if TRT_VERSION_GE(8, 0, 0)
+class Resize2DOpConverter : public TensorRTOpConverter {
+ public:
+ explicit Resize2DOpConverter(std::string op_name)
+ : TensorRTOpConverter(std::move(op_name), {kTensor}) {}
+ ~Resize2DOpConverter() = default;
+
+ void Convert(TensorRTOpConverterParams* params) const {
+ auto input = params->inputs.at(0).tensor;
+ TVM_FFI_ICHECK_EQ(params->node.GetAttr<ffi::String>("layout"), "NCHW");
+ auto input_dims = TrtDimsToVector(input->getDimensions());
+ TVM_FFI_ICHECK_GE(input_dims.size(), 2);
+ // Relax resize2d takes the target (height, width) as a Shape argument
(serialized as arg_size).
+ auto size = params->node.GetAttr<ffi::Array<int64_t>>("arg_size");
+ TVM_FFI_ICHECK_EQ(size.size(), 2);
+ const std::string method = params->node.GetAttr<ffi::String>("method");
+ const std::string coordinate_transformation_mode =
+ params->node.GetAttr<ffi::String>("coordinate_transformation_mode");
+ const std::string rounding_method =
params->node.GetAttr<ffi::String>("rounding_method");
+
+ auto resize_layer = params->network->addResize(*input);
+ TVM_FFI_ICHECK(resize_layer != nullptr);
+
+ static const std::unordered_map<std::string, nvinfer1::InterpolationMode>
method_map = {
+ {"nearest_neighbor", nvinfer1::InterpolationMode::kNEAREST},
+ {"linear", nvinfer1::InterpolationMode::kLINEAR},
+ {"cubic", nvinfer1::InterpolationMode::kCUBIC}};
+ auto method_it = method_map.find(method);
+ TVM_FFI_ICHECK(method_it != method_map.end()) << "Unsupported resize2d
method " << method;
+ resize_layer->setResizeMode(method_it->second);
+
+ // pytorch_half_pixel matches half_pixel for output extents greater than
one. Singleton output
+ // dimensions are handled by the selector below.
+ static const std::unordered_map<std::string,
nvinfer1::ResizeCoordinateTransformation>
+ coordinate_transformation_map = {
+ {"asymmetric",
nvinfer1::ResizeCoordinateTransformation::kASYMMETRIC},
+ {"align_corners",
nvinfer1::ResizeCoordinateTransformation::kALIGN_CORNERS},
+ {"half_pixel",
nvinfer1::ResizeCoordinateTransformation::kHALF_PIXEL},
+ {"pytorch_half_pixel",
nvinfer1::ResizeCoordinateTransformation::kHALF_PIXEL}};
+ auto coordinate_transformation_it =
+ coordinate_transformation_map.find(coordinate_transformation_mode);
+ TVM_FFI_ICHECK(coordinate_transformation_it !=
coordinate_transformation_map.end())
+ << "Unsupported resize2d coordinate_transformation_mode " <<
coordinate_transformation_mode;
+
resize_layer->setCoordinateTransformation(coordinate_transformation_it->second);
+ if (coordinate_transformation_mode == "pytorch_half_pixel") {
+ // PyTorch maps an output dimension of size one to source coordinate
zero, whereas the
+ // regular half-pixel formula selects the center of the input dimension.
+
resize_layer->setSelectorForSinglePixel(nvinfer1::ResizeSelector::kUPPER);
+ }
+
+ if (method == "nearest_neighbor") {
+ static const std::unordered_map<std::string, nvinfer1::ResizeRoundMode>
rounding_map = {
+ {"floor", nvinfer1::ResizeRoundMode::kFLOOR},
+ {"ceil", nvinfer1::ResizeRoundMode::kCEIL},
+ {"round_prefer_ceil", nvinfer1::ResizeRoundMode::kHALF_UP},
+ {"round_prefer_floor", nvinfer1::ResizeRoundMode::kHALF_DOWN}};
+ auto rounding_it = rounding_map.find(rounding_method);
+ TVM_FFI_ICHECK(rounding_it != rounding_map.end())
+ << "Unsupported resize2d rounding_method " << rounding_method;
+ resize_layer->setNearestRounding(rounding_it->second);
+ }
+
+ if (method == "cubic") {
+
resize_layer->setCubicCoeff(static_cast<float>(params->node.GetAttr<double>("cubic_alpha")));
+ resize_layer->setExcludeOutside(
+ static_cast<bool>(params->node.GetAttr<int64_t>("cubic_exclude")));
+ }
+
+ std::vector<int> output_dims(input_dims.begin(), input_dims.end());
+ output_dims[output_dims.size() - 2] = static_cast<int>(size[0]);
+ output_dims[output_dims.size() - 1] = static_cast<int>(size[1]);
+ resize_layer->setOutputDimensions(VectorToTrtDims(output_dims));
+
+ params->outputs.push_back(resize_layer->getOutput(0));
+ }
+};
+#endif // TRT_VERSION_GE(8, 0, 0)
+
class AdaptivePoolingOpConverter : public TensorRTOpConverter {
public:
explicit AdaptivePoolingOpConverter(std::string op_name)
@@ -1291,6 +1386,7 @@ const std::unordered_map<std::string,
std::unique_ptr<TensorRTOpConverter>>& Get
all_converters.emplace_back(std::make_unique<ActivationOpConverter>("nn.relu"));
all_converters.emplace_back(std::make_unique<ActivationOpConverter>("sigmoid"));
all_converters.emplace_back(std::make_unique<ActivationOpConverter>("tanh"));
+ all_converters.emplace_back(std::make_unique<SiluOpConverter>("nn.silu"));
all_converters.emplace_back(std::make_unique<BatchNormOpConverter>("nn.batch_norm"));
all_converters.emplace_back(std::make_unique<LayerNormOpConverter>("nn.layer_norm"));
all_converters.emplace_back(std::make_unique<SoftmaxOpConverter>("nn.softmax"));
@@ -1356,6 +1452,9 @@ const std::unordered_map<std::string,
std::unique_ptr<TensorRTOpConverter>>& Get
#if TRT_VERSION_GE(7, 0, 0)
all_converters.emplace_back(std::make_unique<UnaryOpConverter>("erf"));
#endif // TRT_VERSION_GE(7, 0, 0)
+#if TRT_VERSION_GE(8, 0, 0)
+
all_converters.emplace_back(std::make_unique<Resize2DOpConverter>("image.resize2d"));
+#endif // TRT_VERSION_GE(8, 0, 0)
auto* map = new std::unordered_map<std::string,
std::unique_ptr<TensorRTOpConverter>>();
for (auto& converter : all_converters) {
map->emplace("tensorrt." + converter->op_name, std::move(converter));
diff --git a/tests/python/relax/test_backend_tensorrt.py
b/tests/python/relax/test_backend_tensorrt.py
new file mode 100644
index 0000000000..722481cc0c
--- /dev/null
+++ b/tests/python/relax/test_backend_tensorrt.py
@@ -0,0 +1,117 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pytest
+
+from tvm import relax
+from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
+
+
+def _make_resize2d_module(
+ input_shape=(1, 3, 8, 8),
+ input_dtype="float32",
+ size=(16, 16),
+ *,
+ dynamic_size=False,
+ bind_static_size=False,
+ layout="NCHW",
+ method="linear",
+ coordinate_transformation_mode="half_pixel",
+ rounding_method="round",
+ out_dtype=None,
+):
+ builder = relax.BlockBuilder()
+ data = relax.Var("data", relax.TensorType(input_shape, input_dtype))
+ params = [data]
+ if dynamic_size:
+ size_expr = relax.Var("size", relax.ShapeType(ndim=2))
+ params.append(size_expr)
+ else:
+ size_expr = relax.ShapeExpr(size)
+
+ with builder.function("main", params):
+ if bind_static_size:
+ size_expr = builder.emit(size_expr, "size")
+ with builder.dataflow():
+ output = builder.emit(
+ relax.op.image.resize2d(
+ data,
+ size=size_expr,
+ layout=layout,
+ method=method,
+
coordinate_transformation_mode=coordinate_transformation_mode,
+ rounding_method=rounding_method,
+ out_dtype=out_dtype,
+ )
+ )
+ output = builder.emit_output(output)
+ builder.emit_func_output(output)
+ return builder.get()
+
+
+def _tensorrt_regions(mod):
+ return [
+ func
+ for func in mod.functions.values()
+ if isinstance(func, relax.Function)
+ and func.attrs is not None
+ and func.attrs.get("Codegen") == "tensorrt"
+ ]
+
+
+def test_resize2d_partition_supported():
+ mod = _make_resize2d_module(
+ size=(13, 11),
+ bind_static_size=True,
+ method="nearest_neighbor",
+ coordinate_transformation_mode="asymmetric",
+ rounding_method="floor",
+ )
+ partitioned = partition_for_tensorrt(mod)
+
+ regions = _tensorrt_regions(partitioned)
+ assert len(regions) == 1
+ assert len(regions[0].params) == 1
+
+
[email protected](
+ "kwargs",
+ [
+ pytest.param({"input_shape": (1, 8, 8, 3), "layout": "NHWC"},
id="unsupported-layout"),
+ pytest.param({"dynamic_size": True}, id="dynamic-size"),
+ pytest.param({"input_dtype": "float64"}, id="unsupported-input-dtype"),
+ pytest.param({"out_dtype": "float16"}, id="different-output-dtype"),
+ pytest.param(
+ {
+ "method": "nearest_neighbor",
+ "coordinate_transformation_mode": "tf_half_pixel_for_nn",
+ "rounding_method": "floor",
+ },
+ id="unsupported-coordinate-mode",
+ ),
+ pytest.param(
+ {
+ "method": "nearest_neighbor",
+ "coordinate_transformation_mode": "asymmetric",
+ "rounding_method": "round",
+ },
+ id="ties-to-even-rounding",
+ ),
+ ],
+)
+def test_resize2d_partition_fallback(kwargs):
+ partitioned = partition_for_tensorrt(_make_resize2d_module(**kwargs))
+ assert not _tensorrt_regions(partitioned)
diff --git a/tests/python/relax/test_codegen_tensorrt.py
b/tests/python/relax/test_codegen_tensorrt.py
index 73e3471b20..4a3c35ebd7 100644
--- a/tests/python/relax/test_codegen_tensorrt.py
+++ b/tests/python/relax/test_codegen_tensorrt.py
@@ -21,6 +21,7 @@ import tvm
import tvm.testing
from tvm import relax
from tvm.contrib.pickle_memoize import memoize
+from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
from tvm.relax.dpl import is_op, make_fused_bias_activation_pattern, wildcard
from tvm.script import relax as R
from tvm.testing import env
@@ -128,13 +129,13 @@ def _offload_and_compare(mod, params_np, patterns,
data_np, rtol=1e-2, atol=1e-2
otherwise collapse repeated ops.
"""
ref = build_and_run(mod, [data_np, *params_np.values()], "llvm",
legalize=True)
- partitioned = tvm.transform.Sequential(
- [
- relax.transform.BindParams("main", params_np),
- relax.transform.FuseOpsByPattern(patterns),
- relax.transform.MergeCompositeFunctions(),
- ]
- )(mod)
+ mod = relax.transform.BindParams("main", params_np)(mod)
+ if patterns is None:
+ partitioned = partition_for_tensorrt(mod)
+ else:
+ partitioned = tvm.transform.Sequential(
+ [relax.transform.FuseOpsByPattern(patterns),
relax.transform.MergeCompositeFunctions()]
+ )(mod)
# Guard against a silent false pass: if no pattern matched, nothing is
offloaded and the
# comparison would trivially succeed via the TVM fallback without
exercising the converter.
assert any(
@@ -224,6 +225,26 @@ def test_tensorrt_sigmoid():
_offload_and_compare(Sigmoid, {}, patterns, data)
+def test_tensorrt_subtract_constant_lhs():
+ """A bound constant on the left-hand side must remain the minuend."""
+
+ @tvm.script.ir_module
+ class ConstantMinusTensor:
+ @R.function
+ def main(
+ data: R.Tensor((2, 3, 4), "float32"),
+ constant: R.Tensor((2, 3, 4), "float32"),
+ ):
+ with R.dataflow():
+ out = relax.op.subtract(constant, data)
+ R.output(out)
+ return out
+
+ data = np.linspace(-2.0, 3.0, 24, dtype="float32").reshape(2, 3, 4)
+ constant = np.linspace(4.0, 9.0, 24, dtype="float32").reshape(2, 3, 4)
+ _offload_and_compare(ConstantMinusTensor, {"constant": constant}, None,
data)
+
+
def test_tensorrt_tanh():
@tvm.script.ir_module
class Tanh:
@@ -475,18 +496,13 @@ def test_tensorrt_split():
def main(data: R.Tensor((4, 8, 16), "float32")):
with R.dataflow():
parts = relax.op.split(data, 2, axis=1)
- out = relax.op.add(parts[0], parts[1])
+ out = relax.op.subtract(parts[1], parts[0])
R.output(out)
return out
- data = np.random.randn(4, 8, 16).astype("float32")
- # Offload the add too so both split outputs are consumed inside TensorRT
(and nothing is left
- # for the VM to legalize).
- patterns = [
- ("tensorrt.split", is_op("relax.split")(wildcard())),
- ("tensorrt.add", is_op("relax.add")(wildcard(), wildcard())),
- ]
- _offload_and_compare(Split, {}, patterns, data)
+ data = np.ones((4, 8, 16), dtype="float32")
+ data[:, 4:, :] = 5
+ _offload_and_compare(Split, {}, None, data)
def test_tensorrt_layout_transform():
@@ -640,5 +656,78 @@ def test_partition_for_tensorrt():
tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)
[email protected](
+ "out_hw, method, coordinate_transformation_mode, rounding_method,
cubic_alpha",
+ [
+ ((13, 11), "nearest_neighbor", "asymmetric", "floor", -0.75),
+ ((13, 11), "linear", "half_pixel", "round", -0.75),
+ ((16, 16), "cubic", "half_pixel", "round", -0.5),
+ ],
+)
+def test_tensorrt_resize2d(
+ out_hw, method, coordinate_transformation_mode, rounding_method,
cubic_alpha
+):
+ """Regression test for image.resize2d offload in #19887."""
+
+ @tvm.script.ir_module
+ class Resize2D:
+ @R.function
+ def main(data: R.Tensor((1, 3, 8, 8), "float32")):
+ with R.dataflow():
+ out = relax.op.image.resize2d(
+ data,
+ size=out_hw,
+ layout="NCHW",
+ method=method,
+
coordinate_transformation_mode=coordinate_transformation_mode,
+ rounding_method=rounding_method,
+ cubic_alpha=cubic_alpha,
+ )
+ R.output(out)
+ return out
+
+ data = np.random.randn(1, 3, 8, 8).astype("float32")
+ _offload_and_compare(Resize2D, {}, None, data)
+
+
+def test_tensorrt_resize2d_pytorch_half_pixel_single_dimension():
+ """pytorch_half_pixel selects source coordinate zero for each singleton
output dimension."""
+
+ @tvm.script.ir_module
+ class ResizeSingleDimension:
+ @R.function
+ def main(data: R.Tensor((1, 1, 5, 7), "float32")):
+ with R.dataflow():
+ out = relax.op.image.resize2d(
+ data,
+ size=(1, 1),
+ method="nearest_neighbor",
+ coordinate_transformation_mode="pytorch_half_pixel",
+ rounding_method="floor",
+ )
+ R.output(out)
+ return out
+
+ mod = ResizeSingleDimension
+ data = np.arange(35, dtype="float32").reshape(1, 1, 5, 7)
+ _offload_and_compare(mod, {}, None, data, rtol=0, atol=0)
+
+
+def test_tensorrt_silu():
+ """YOLO's SiLU activation is lowered as x * sigmoid(x)."""
+
+ @tvm.script.ir_module
+ class Silu:
+ @R.function
+ def main(data: R.Tensor((1, 16, 8, 8), "float32")):
+ with R.dataflow():
+ out = relax.op.nn.silu(data)
+ R.output(out)
+ return out
+
+ data = np.random.randn(1, 16, 8, 8).astype("float32")
+ _offload_and_compare(Silu, {}, None, data)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_op_create.py
b/tests/python/relax/test_op_create.py
index c9981de37a..0e8d403bae 100644
--- a/tests/python/relax/test_op_create.py
+++ b/tests/python/relax/test_op_create.py
@@ -567,6 +567,11 @@ def test_arange_infer_ty():
_check_inference(bb, relax.op.arange(1.0, 10), relax.TensorType((9,),
"float32"))
_check_inference(bb, relax.op.arange(0, 20, 2.5), relax.TensorType((8,),
"float32"))
_check_inference(bb, relax.op.arange(1, 10, 2.3), relax.TensorType((4,),
"float32"))
+ _check_inference(
+ bb,
+ relax.op.arange(0, 10, 1, tvm.ir.PrimType("float32")),
+ relax.TensorType((10,), "float32"),
+ )
def test_arange_infer_ty_shape_var():
diff --git a/tests/python/relax/test_transform.py
b/tests/python/relax/test_transform.py
index d32caa69d5..dd811855d8 100644
--- a/tests/python/relax/test_transform.py
+++ b/tests/python/relax/test_transform.py
@@ -356,6 +356,50 @@ def test_call_dps_packed_rewrite():
assert s2.op.global_symbol == "test.op.identity"
+def test_call_dps_packed_rewrite_nested_tuple_output():
+ """Flatten nested outputs for the packed ABI, then rebuild their Relax
structure."""
+ input_ty = relax.TensorType((2, 3), "float32")
+ flat_output_types = [
+ relax.TensorType((2, 3), "float32"),
+ relax.TensorType((4,), "int32"),
+ relax.TensorType((5, 6), "float16"),
+ ]
+ output_ty = tvm.ir.TupleType([flat_output_types[0],
tvm.ir.TupleType(flat_output_types[1:])])
+
+ x = relax.Var("x", input_ty)
+ call = relax.Call(
+ tvm.ir.Op.get("relax.call_dps_packed"),
+ [relax.ExternFunc("test.op.nested_outputs"), relax.Tuple([x])],
+ ty_args=[output_ty],
+ )
+ builder = relax.BlockBuilder()
+ with builder.function("main", [x], attrs={"relax.force_pure": True}):
+ out = builder.emit(call)
+ builder.emit_func_output(out)
+
+ after = relax.transform.CallTIRRewrite()(builder.get())
+ relax.analysis.well_formed(after)
+ func = after["main"]
+ block = func.body.blocks[0]
+
+ alloc_bindings = block.bindings[:3]
+ for binding, expected_ty in zip(alloc_bindings, flat_output_types):
+ assert binding.value.op.name == "relax.builtin.alloc_tensor"
+ tvm.ir.assert_structural_equal(binding.var.ty, expected_ty)
+
+ packed_call = block.bindings[3].value
+ assert packed_call.op.global_symbol == "test.op.nested_outputs"
+ assert packed_call.args[0].same_as(func.params[0])
+ assert all(
+ arg.same_as(binding.var) for arg, binding in zip(packed_call.args[1:],
alloc_bindings)
+ )
+
+ rebuilt = block.bindings[4].value
+ assert rebuilt.fields[0].same_as(alloc_bindings[0].var)
+ assert rebuilt.fields[1].fields[0].same_as(alloc_bindings[1].var)
+ assert rebuilt.fields[1].fields[1].same_as(alloc_bindings[2].var)
+
+
def test_call_tir_inplace_simple():
# simple case: one inplace argument
@tvm.script.ir_module
diff --git a/tests/python/relax/test_transform_fuse_ops_by_pattern.py
b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
index 05102dd628..52e5f40d47 100644
--- a/tests/python/relax/test_transform_fuse_ops_by_pattern.py
+++ b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
@@ -1422,5 +1422,81 @@ def test_concat():
check(mod, [("x.concat", pat_clip)], Expected2)
+def test_unique_boundary_output_precedes_last_group_binding():
+ """Export a sole boundary output even when a dead internal binding follows
it."""
+
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(
+ x: R.Tensor((2, 4), "float32"),
+ ) -> R.Tensor((2, 4), "float32"):
+ with R.dataflow():
+ first = R.nn.relu(x)
+ kept = R.nn.relu(first)
+ dead = R.nn.relu(kept)
+ R.output(kept)
+ return kept
+
+ pattern =
is_op("relax.nn.relu")(is_op("relax.nn.relu")(is_op("relax.nn.relu")(wildcard())))
+ after = relax.transform.FuseOpsByPattern(
+ [("compiler_A.relu_chain", pattern)], annotate_codegen=True
+ )(Before)
+
+ relax.analysis.well_formed(after)
+ assert not relax.analysis.free_vars(after["main"])
+
+ calls = [
+ binding
+ for block in after["main"].body.blocks
+ for binding in block.bindings
+ if isinstance(binding.value, relax.Call)
+ and isinstance(binding.value.op, relax.GlobalVar)
+ and after[binding.value.op].attrs.get("Codegen") == "compiler_A"
+ ]
+ assert len(calls) == 1
+ grouped_result = calls[0].var
+ assert not isinstance(grouped_result, relax.DataflowVar)
+ assert after["main"].body.body.same_as(grouped_result)
+
+
+def test_inline_bound_static_shape_argument():
+ """A static leaf binding should not become a grouped-function parameter."""
+
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor((4,), "float32")) -> R.Tensor((2, 2), "float32"):
+ shape: R.Shape([2, 2]) = R.shape([2, 2])
+ with R.dataflow():
+ out: R.Tensor((2, 2), "float32") = R.reshape(x, shape)
+ R.output(out)
+ return out
+
+ pattern = is_op("relax.reshape")(wildcard(), wildcard())
+ after = relax.transform.FuseOpsByPattern([("compiler_A.reshape",
pattern)])(Before)
+ grouped = [
+ func
+ for func in after.functions.values()
+ if isinstance(func, relax.Function)
+ and func.attrs is not None
+ and func.attrs.get("Composite") == "compiler_A.reshape"
+ ]
+ assert len(grouped) == 1
+ assert len(grouped[0].params) == 1
+
+ reshape_calls = []
+ relax.analysis.post_order_visit(
+ grouped[0],
+ lambda expr: reshape_calls.append(expr)
+ if isinstance(expr, relax.Call)
+ and isinstance(expr.op, tvm.ir.Op)
+ and expr.op.name == "relax.reshape"
+ else None,
+ )
+ assert len(reshape_calls) == 1
+ assert isinstance(reshape_calls[0].args[1], relax.ShapeExpr)
+
+
if __name__ == "__main__":
pytest.main([__file__])
diff --git a/tests/python/relax/test_transform_merge_composite_functions.py
b/tests/python/relax/test_transform_merge_composite_functions.py
index b1ce09799a..86019e706e 100644
--- a/tests/python/relax/test_transform_merge_composite_functions.py
+++ b/tests/python/relax/test_transform_merge_composite_functions.py
@@ -931,6 +931,16 @@ def check(mod, expected):
tvm.ir.assert_structural_equal(partitioned, expected)
+def get_codegen_regions(mod):
+ return {
+ gvar.name_hint: str(func.attrs["Codegen"])
+ for gvar, func in mod.functions.items()
+ if isinstance(func, relax.Function)
+ and func.attrs is not None
+ and func.attrs.get("Codegen") is not None
+ }
+
+
def test_conv2d_relu_x2():
check(Conv2dReLUx2, Conv2dReLUx2_merged)
@@ -1221,5 +1231,181 @@ def test_handle_existence_of_call_tir():
tvm.ir.assert_structural_equal(Expected, After)
+def test_tuple_projection_merging():
+ """Merge projections without changing their indices; visit inline tuple
arguments."""
+
+ @tvm.script.ir_module
+ class Before:
+ @R.function(private=True)
+ def split(
+ x: R.Tensor((2, 4), "float32"),
+ ) -> R.Tuple(R.Tensor((1, 4), "float32"), R.Tensor((1, 4), "float32")):
+ R.func_attr({"Composite": "compiler_A.split", "Primitive": True})
+ return R.split(x, indices_or_sections=2, axis=0)
+
+ @R.function(private=True)
+ def subtract(
+ x: R.Tensor((1, 4), "float32"),
+ y: R.Tensor((1, 4), "float32"),
+ ) -> R.Tensor((1, 4), "float32"):
+ R.func_attr({"Composite": "compiler_A.subtract", "Primitive":
True})
+ return R.subtract(x, y)
+
+ @R.function
+ def main(x: R.Tensor((2, 4), "float32")) -> R.Tensor((2, 4),
"float32"):
+ cls = Before
+ with R.dataflow():
+ parts = cls.split(x)
+ left = parts[0]
+ right = parts[1]
+ repacked = R.tuple(left, right)
+ reordered_left = repacked[0]
+ reordered_right = repacked[1]
+ difference = cls.subtract(reordered_right, reordered_left)
+ out = R.concat((difference, difference), axis=0)
+ R.output(out)
+ return out
+
+ after = relax.transform.MergeCompositeFunctions()(Before)
+ assert get_codegen_regions(after) == {"fused_split_subtract_compiler_A":
"compiler_A"}
+ relax.analysis.well_formed(after)
+
+ region = next(
+ func
+ for func in after.functions.values()
+ if isinstance(func, relax.Function)
+ and func.attrs is not None
+ and func.attrs.get("Codegen") == "compiler_A"
+ )
+ bindings = [binding for block in region.body.blocks for binding in
block.bindings]
+ bound_values = {binding.var: binding.value for binding in bindings}
+ subtract_var = next(
+ var
+ for var, value in bound_values.items()
+ if isinstance(value, relax.Function)
+ and value.attrs is not None
+ and value.attrs.get("Composite") == "compiler_A.subtract"
+ )
+ subtract_call = next(
+ value
+ for value in bound_values.values()
+ if isinstance(value, relax.Call) and value.op.same_as(subtract_var)
+ )
+ assert [bound_values[arg].index for arg in subtract_call.args] == [1, 0]
+
+
+def test_tuple_projection_preserves_real_cycle_boundary():
+ """Transparent projections must not bridge a dependency through another
codegen."""
+
+ @tvm.script.ir_module
+ class Before:
+ @R.function(private=True)
+ def split(
+ x: R.Tensor((2, 4), "float32"),
+ ) -> R.Tuple(R.Tensor((1, 4), "float32"), R.Tensor((1, 4), "float32")):
+ R.func_attr({"Composite": "compiler_A.split", "Primitive": True})
+ return R.split(x, indices_or_sections=2, axis=0)
+
+ @R.function(private=True)
+ def foreign_relu(
+ x: R.Tensor((1, 4), "float32"),
+ ) -> R.Tensor((1, 4), "float32"):
+ R.func_attr({"Composite": "compiler_B.relu", "Primitive": True})
+ return R.nn.relu(x)
+
+ @R.function(private=True)
+ def add(
+ x: R.Tensor((1, 4), "float32"),
+ y: R.Tensor((1, 4), "float32"),
+ ) -> R.Tensor((1, 4), "float32"):
+ R.func_attr({"Composite": "compiler_A.add", "Primitive": True})
+ return R.add(x, y)
+
+ @R.function
+ def main(x: R.Tensor((2, 4), "float32")) -> R.Tensor((1, 4),
"float32"):
+ cls = Before
+ with R.dataflow():
+ parts = cls.split(x)
+ left = parts[0]
+ right = parts[1]
+ foreign = cls.foreign_relu(left)
+ out = cls.add(right, foreign)
+ R.output(out)
+ return out
+
+ after = relax.transform.MergeCompositeFunctions()(Before)
+ assert get_codegen_regions(after) == {
+ "fused_add_compiler_A": "compiler_A",
+ "fused_foreign_relu_compiler_B": "compiler_B",
+ "fused_split_compiler_A": "compiler_A",
+ }
+ relax.analysis.well_formed(after)
+
+
+def test_tuple_projection_rejects_escaping_tuple():
+ """Keep the projection outside when the complete tuple also crosses the
region boundary."""
+
+ @tvm.script.ir_module
+ class Before:
+ @R.function(private=True)
+ def split(
+ x: R.Tensor((2, 4), "float32"),
+ ) -> R.Tuple(R.Tensor((1, 4), "float32"), R.Tensor((1, 4), "float32")):
+ R.func_attr({"Composite": "compiler_A.split", "Primitive": True})
+ return R.split(x, indices_or_sections=2, axis=0)
+
+ @R.function(private=True)
+ def relu(
+ x: R.Tensor((1, 4), "float32"),
+ ) -> R.Tensor((1, 4), "float32"):
+ R.func_attr({"Composite": "compiler_A.relu", "Primitive": True})
+ return R.nn.relu(x)
+
+ @R.function
+ def main(
+ x: R.Tensor((2, 4), "float32"),
+ ) -> R.Tuple(
+ R.Tuple(R.Tensor((1, 4), "float32"), R.Tensor((1, 4), "float32")),
+ R.Tensor((1, 4), "float32"),
+ ):
+ cls = Before
+ with R.dataflow():
+ parts = cls.split(x)
+ alias = parts
+ left = alias[0]
+ out = cls.relu(left)
+ result = R.tuple(alias, out)
+ R.output(result)
+ return result
+
+ after = relax.transform.MergeCompositeFunctions()(Before)
+ assert get_codegen_regions(after) == {
+ "fused_relu_compiler_A": "compiler_A",
+ "fused_split_compiler_A": "compiler_A",
+ }
+ relax.analysis.well_formed(after)
+
+ bindings = [binding for block in after["main"].body.blocks for binding in
block.bindings]
+ split_result = next(
+ binding.var
+ for binding in bindings
+ if isinstance(binding.value, relax.Call)
+ and isinstance(binding.value.op, relax.GlobalVar)
+ and binding.value.op.name_hint == "fused_split_compiler_A"
+ )
+ projections = [
+ binding.value
+ for binding in bindings
+ if isinstance(binding.value, relax.TupleGetItem)
+ and binding.value.tuple_value.same_as(split_result)
+ ]
+ assert [projection.index for projection in projections] == [0]
+
+ result_var = after["main"].body.body
+ result_value = next(binding.value for binding in bindings if
binding.var.same_as(result_var))
+ assert isinstance(result_value, relax.Tuple)
+ assert result_value.fields[0].same_as(split_result)
+
+
if __name__ == "__main__":
pytest.main([__file__])
