diff --git a/python/tvm/relax/backend/contrib/tensorrt.py b/python/tvm/relax/backend/contrib/tensorrt.py index 303ebc394cb9..08be7dc1a36b 100644 --- a/python/tvm/relax/backend/contrib/tensorrt.py +++ b/python/tvm/relax/backend/contrib/tensorrt.py @@ -47,6 +47,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 +93,9 @@ def _tensorrt_patterns() -> list[Pattern]: ]: patterns.append(_op_pattern(composite, op, 2)) + # image.resize2d (data + target-size shape argument). + patterns.append(_op_pattern("tensorrt.image.resize2d", "relax.image.resize2d", 2)) + # 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 e8c1300f012c..ace1e13258cb 100644 --- a/python/tvm/relax/op/create.py +++ b/python/tvm/relax/op/create.py @@ -288,7 +288,7 @@ def is_int(expr): 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/python/tvm/s_tir/dlight/gpu/gemv.py b/python/tvm/s_tir/dlight/gpu/gemv.py index c893c3449a95..1c451b964faa 100644 --- a/python/tvm/s_tir/dlight/gpu/gemv.py +++ b/python/tvm/s_tir/dlight/gpu/gemv.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# ruff: noqa: E741, F821 +# ruff: noqa: E741 """A rule for GEMV and DecodeGEMV.""" from functools import reduce diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py index e854ec1e2639..8dedbec4b661 100644 --- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py +++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# ruff: noqa: E741, F821 +# ruff: noqa: E741 """A rule for low-batch GEMM / decode-GEMM using GEMV schedule.""" from functools import reduce diff --git a/src/relax/backend/contrib/tensorrt/codegen.cc b/src/relax/backend/contrib/tensorrt/codegen.cc index c5338e2d2732..bd088c4022ba 100644 --- a/src/relax/backend/contrib/tensorrt/codegen.cc +++ b/src/relax/backend/contrib/tensorrt/codegen.cc @@ -26,14 +26,17 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include +#include #include #include #include @@ -215,6 +218,12 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { TensorRTJSONSerializer* serializer_; /*! \brief Accumulated translated arguments. */ std::vector args_; + /*! \brief The first primitive operator call in the composite function. */ + const CallNode* operator_call_{nullptr}; + /*! \brief Total number of calls in the composite function. */ + size_t num_calls_{0}; + /*! \brief Number of primitive operator calls in the composite function. */ + size_t num_operator_calls_{0}; /*! * \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. @@ -249,16 +258,74 @@ class TensorRTJSONSerializer : public JSONSerializer { CollectFromCompositeFunctionBody collector(this); collector.VisitExpr(fn->body); - // Capture the args to the "Composite" function as inputs for this node. - std::vector inputs; - for (const auto& arg : call_node->args) { + // Capture the args to the "Composite" function as inputs for this node, and bind the + // composite parameters back to those caller-side expressions. + TVM_FFI_ICHECK_EQ(fn->params.size(), call_node->args.size()); + ffi::Map param_bindings; + std::vector arguments; + for (size_t i = 0; i < call_node->args.size(); ++i) { + const Expr& arg = call_node->args[i]; + param_bindings.Set(fn->params[i], arg); auto res = VisitExpr(arg); - inputs.insert(inputs.end(), res.begin(), res.end()); + arguments.insert(arguments.end(), res.begin(), res.end()); } // Capture constants from the composite function body as additional inputs for this node. - for (const auto& node : collector.args_) { - inputs.emplace_back(node); + // FuseOps keeps constants in the composite body, so the legacy serializer discovered caller + // arguments and constants separately and always placed constants last. For the single-op + // composites registered by TensorRT, bind the operator arguments back to caller expressions + // and visit them in their original order. This matters for non-commutative operators such as + // subtract(constant, tensor), as well as for constants interleaved in tuple arguments. + std::vector inputs; + bool inputs_are_ordered = collector.num_calls_ == 1 && collector.num_operator_calls_ == 1; + if (inputs_are_ordered) { + TVM_FFI_ICHECK_NOTNULL(collector.operator_call_); + auto is_resolvable = [&](const auto& self, const Expr& expr) -> bool { + if (const auto* tuple = expr.as()) { + return std::all_of(tuple->fields.begin(), tuple->fields.end(), + [&](const Expr& field) { return self(self, field); }); + } + Type type = GetType(expr); + if (!type->IsInstance() && !type->IsInstance()) return true; + if (expr->IsInstance()) return true; + if (const auto* var = expr.as()) { + return param_bindings.count(ffi::GetRef(var)); + } + if (expr->IsInstance()) { + for (const Var& var : FreeVars(expr)) { + if (!param_bindings.count(var)) return false; + } + return true; + } + return false; + }; + for (const Expr& arg : collector.operator_call_->args) { + if (!is_resolvable(is_resolvable, arg)) { + inputs_are_ordered = false; + break; + } + } + } + if (inputs_are_ordered) { + auto append_tensor_inputs = [&](const auto& self, const Expr& expr) -> void { + if (const auto* tuple = expr.as()) { + for (const Expr& field : tuple->fields) self(self, field); + return; + } + Type type = GetType(expr); + if (type->IsInstance() || type->IsInstance()) { + 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)); + } + } else { + VLOG(1) << name << " is not a directly resolvable single-op composite; preserving the " + << "legacy argument ordering"; + inputs = std::move(arguments); + inputs.insert(inputs.end(), collector.args_.begin(), collector.args_.end()); } // Create the final node. @@ -306,6 +373,11 @@ void CollectFromCompositeFunctionBody::VisitExpr_(const ConstantNode* constant_n } void CollectFromCompositeFunctionBody::VisitExpr_(const CallNode* call_node) { + ++num_calls_; + if (call_node->op->IsInstance()) { + ++num_operator_calls_; + if (operator_call_ == nullptr) 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 8c1a16c66ade..f409ade79c1f 100644 --- a/src/relax/transform/call_tir_rewrite.cc +++ b/src/relax/transform/call_tir_rewrite.cc @@ -29,6 +29,8 @@ #include #include +#include + #include "utils.h" namespace tvm { @@ -75,27 +77,12 @@ class CallTIRMutator : public ExprMutator { bool is_inplace = call->op.same_as(call_tir_inplace_op); const auto* inplace_attrs = call->attrs.as(); ffi::Array outs; + ffi::Optional tuple_output_type; if (const auto& tensor_ty = MatchType(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(), - 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) @@ -107,35 +94,30 @@ class CallTIRMutator : public ExprMutator { } else if (const auto& tuple_ty = MatchType(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()) - << "call_tir expects Tuple of TensorType, but got " << field - << " as an element of TupleType"; - const auto& field_tensor = field.as_or_throw(); - 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(), - 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()->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(); }); + + 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()) + << "call_tir expects Tuple of TensorType, but got " << field + << " as an element of TupleType"; + const auto& field_tensor = field.as_or_throw(); + + 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()->fields[inplace_attrs->inplace_indices[i]]); + } } } } else { @@ -175,15 +157,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(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(), + 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* outs) { + if (const auto* tensor_ty = type.as()) { + outs->push_back(AllocateOutputTensor(ffi::GetRef(tensor_ty), alloc_tensor_op)); + return; + } + if (const auto* tuple_ty = type.as()) { + 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& outs, size_t* index) { + if (type->IsInstance()) { + TVM_FFI_ICHECK_LT(*index, outs.size()); + return outs[(*index)++]; + } + if (const auto* tuple_ty = type.as()) { + ffi::Array 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 15d84930920e..5c302963500f 100644 --- a/src/relax/transform/fuse_ops.cc +++ b/src/relax/transform/fuse_ops.cc @@ -474,16 +474,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& 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. @@ -769,22 +768,6 @@ class OperatorFusor : public ExprMutator { return obj2group; } - bool IsTupleOutput(Function f) { - auto ty = GetType(f).as(); - TVM_FFI_ICHECK(ty); - return ty->ret->IsInstance(); - } - - bool IsNestedTupleOutput(Function f) { - if (!IsTupleOutput(f)) return false; - - auto tup = GetType(f).as()->ret.as(); - for (const auto& field : tup->fields) { - if (field->IsInstance()) return true; - } - return false; - } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) final { group2func_.clear(); @@ -807,9 +790,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> 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> 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 @@ -836,15 +819,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() || 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. @@ -862,29 +840,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()) { - new_var = builder_->Emit(call_to_emit); - } else { + if (output_vars.size() == 1 && !output_vars[0]->IsInstance()) { new_var = builder_->EmitOutput(call_to_emit); + } else { + new_var = builder_->Emit(call_to_emit); + } + + // 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(output_vars[0])] = new_var; + continue; } - // 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); + std::unordered_set 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(std::distance(output_vars.begin(), it)); + TupleGetItem tuple_get(new_var, index); + if (output_var->IsInstance()) { + 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(output_var)); } - } else { - var_remap_[var_binding->var] = new_var; } } // Step 5. Finish the binding block generation. @@ -939,8 +938,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); } } }; @@ -1039,9 +1037,6 @@ class OperatorFusor : public ExprMutator { GroupMap obj2group_; /*! \brief Internal function information map. */ std::unordered_map group2func_; - /*! \brief Record the index for TupleGetItem if the variable needs to be remapped to an output - * tuple element after fusion. */ - std::unordered_map tuple_get_indices_; /*! * \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 18e2736ecb2f..ef1cb943a686 100644 --- a/src/relax/transform/merge_composite_functions.cc +++ b/src/relax/transform/merge_composite_functions.cc @@ -89,7 +89,6 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { // Make default groups for dataflow nodes other than CallNode. // Groups for CallNode are created in its visitor. if (e->IsInstance() || e->IsInstance() || - e->IsInstance() || e->IsInstance() || (!e->IsInstance() && !e->IsInstance() && e.as())) { memo_[e] = arena_->make(); } @@ -166,6 +165,18 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { return group; } + Group* VisitExpr_(const TupleNode* tuple) { + Group* group = arena_->make(); + UpdateGroupDependencies(group, tuple->fields); + return group; + } + + Group* VisitExpr_(const TupleGetItemNode* tuple_get_item) { + Group* group = arena_->make(); + UpdateGroupDependencies(group, {tuple_get_item->tuple}); + return group; + } + private: ffi::Optional GetCodegenName(const Expr& callee) { auto const* gvar = callee.as(); diff --git a/src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc b/src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc index 10b3bdf44760..94f05a9402f8 100644 --- a/src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc +++ b/src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc @@ -27,8 +27,10 @@ #include #include +#include #include #include +#include #include "tensorrt_logger.h" #include "tensorrt_ops.h" @@ -109,13 +111,35 @@ void TensorRTBuilder::AddOutput(const JSONGraphNodeEntry& node, uint32_t entry_i entry_id_map_[name] = entry_id; } +namespace { +std::string SupportedOperatorList( + const std::unordered_map>& map) { + std::vector names; + names.reserve(map.size()); + for (const auto& kv : map) names.push_back(kv.first); + std::sort(names.begin(), names.end()); + std::string out; + for (size_t i = 0; i < names.size(); ++i) { + if (i) out += ", "; + out += names[i]; + } + return out; +} +} // namespace + void TensorRTBuilder::AddLayer(int nid, const JSONGraphNode& node) { TensorRTOpConverterParams params(network_, nid, node, &trt_weights_); // Look up converter. const std::unordered_map>& map = GetOpConverters(); auto it = map.find(params.op_name); - TVM_FFI_ICHECK(it != map.end()) << params.op_name << ": Unsupported operator"; + TVM_FFI_ICHECK(it != map.end()) + << params.op_name + << ": Unsupported operator for the TensorRT BYOC backend. The composite function name must " + "match a registered TensorRT converter; prefer " + "tvm.relax.backend.contrib.tensorrt.partition_for_tensorrt over hand-written patterns. " + "Supported operators: " + << SupportedOperatorList(map); const TensorRTOpConverter& converter = *it->second; if (!converter.variable_input_count) { TVM_FFI_ICHECK_EQ(node.GetInputs().size(), converter.input_types.size()) diff --git a/src/runtime/extra/contrib/tensorrt/tensorrt_ops.cc b/src/runtime/extra/contrib/tensorrt/tensorrt_ops.cc index 61dee5a0fbb9..02b5f35fba6e 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,80 @@ 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("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>("arg_size"); + TVM_FFI_ICHECK_EQ(size.size(), 2); + const std::string method = params->node.GetAttr("method"); + const std::string coordinate_transformation_mode = + params->node.GetAttr("coordinate_transformation_mode"); + const std::string rounding_method = params->node.GetAttr("rounding_method"); + + auto resize_layer = params->network->addResize(*input); + TVM_FFI_ICHECK(resize_layer != nullptr); + + static const std::unordered_map 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 every output extent > 1; BYOC shapes are static so + // the resized height/width are concrete and never 1 for an actual upsample. + static const std::unordered_map + 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 (method == "nearest_neighbor") { + static const std::unordered_map rounding_map = { + {"floor", nvinfer1::ResizeRoundMode::kFLOOR}, + {"ceil", nvinfer1::ResizeRoundMode::kCEIL}, + {"round", nvinfer1::ResizeRoundMode::kHALF_UP}, + {"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(params->node.GetAttr("cubic_alpha"))); + resize_layer->setExcludeOutside( + static_cast(params->node.GetAttr("cubic_exclude"))); + } + + std::vector output_dims(input_dims.begin(), input_dims.end()); + output_dims[output_dims.size() - 2] = static_cast(size[0]); + output_dims[output_dims.size() - 1] = static_cast(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 +1382,7 @@ const std::unordered_map>& Get all_converters.emplace_back(std::make_unique("nn.relu")); all_converters.emplace_back(std::make_unique("sigmoid")); all_converters.emplace_back(std::make_unique("tanh")); + all_converters.emplace_back(std::make_unique("nn.silu")); all_converters.emplace_back(std::make_unique("nn.batch_norm")); all_converters.emplace_back(std::make_unique("nn.layer_norm")); all_converters.emplace_back(std::make_unique("nn.softmax")); @@ -1356,6 +1448,9 @@ const std::unordered_map>& Get #if TRT_VERSION_GE(7, 0, 0) all_converters.emplace_back(std::make_unique("erf")); #endif // TRT_VERSION_GE(7, 0, 0) +#if TRT_VERSION_GE(8, 0, 0) + all_converters.emplace_back(std::make_unique("image.resize2d")); +#endif // TRT_VERSION_GE(8, 0, 0) auto* map = new std::unordered_map>(); for (auto& converter : all_converters) { map->emplace("tensorrt." + converter->op_name, std::move(converter)); diff --git a/tests/python/relax/test_codegen_tensorrt.py b/tests/python/relax/test_codegen_tensorrt.py index f6f93aca4d26..01cdcb33c15e 100644 --- a/tests/python/relax/test_codegen_tensorrt.py +++ b/tests/python/relax/test_codegen_tensorrt.py @@ -224,6 +224,48 @@ def main(data: R.Tensor((2, 8, 16, 16), "float32")): _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) + patterns = [("tensorrt.subtract", is_op("relax.subtract")(wildcard(), wildcard()))] + _offload_and_compare(ConstantMinusTensor, {"constant": constant}, patterns, data) + + +def test_tensorrt_subtract_constant_rhs(): + """A bound constant on the right-hand side must remain the subtrahend.""" + + @tvm.script.ir_module + class TensorMinusConstant: + @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(data, constant) + 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) + patterns = [("tensorrt.subtract", is_op("relax.subtract")(wildcard(), wildcard()))] + _offload_and_compare(TensorMinusConstant, {"constant": constant}, patterns, data) + + def test_tensorrt_tanh(): @tvm.script.ir_module class Tanh: @@ -640,5 +682,133 @@ def main( tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2) +def test_tensorrt_permute_dims(): + """Regression test for the composite/runtime converter naming mismatch in #19887.""" + + @tvm.script.ir_module + class PermuteDims: + @R.function + def main(data: R.Tensor((1, 3, 4, 5), "float32")): + with R.dataflow(): + out = relax.op.permute_dims(data, axes=[0, 2, 3, 1]) + R.output(out) + return out + + data = np.random.randn(1, 3, 4, 5).astype("float32") + patterns = [("tensorrt.transpose", is_op("relax.permute_dims")(wildcard()))] + _offload_and_compare(PermuteDims, {}, patterns, data) + + +@pytest.mark.parametrize( + "out_hw, method, coordinate_transformation_mode, rounding_method", + [ + ((16, 16), "nearest_neighbor", "asymmetric", "floor"), + ((16, 16), "linear", "half_pixel", "round"), + ((13, 11), "nearest_neighbor", "asymmetric", "floor"), + ((13, 11), "linear", "half_pixel", "round"), + ((5, 6), "linear", "half_pixel", "round"), + ((16, 16), "linear", "align_corners", "round"), + ((13, 11), "nearest_neighbor", "asymmetric", "round_prefer_floor"), + ((13, 11), "linear", "pytorch_half_pixel", "round"), + ], +) +def test_tensorrt_resize2d(out_hw, method, coordinate_transformation_mode, rounding_method): + """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, + ) + R.output(out) + return out + + data = np.random.randn(1, 3, 8, 8).astype("float32") + patterns = [("tensorrt.image.resize2d", is_op("relax.image.resize2d")(wildcard(), wildcard()))] + _offload_and_compare(Resize2D, {}, patterns, data) + + +def test_tensorrt_resize2d_cubic(): + @tvm.script.ir_module + class Cubic: + @R.function + def main(data: R.Tensor((1, 3, 8, 8), "float32")): + with R.dataflow(): + out = relax.op.image.resize2d( + data, + size=(16, 16), + layout="NCHW", + method="cubic", + coordinate_transformation_mode="half_pixel", + cubic_alpha=-0.5, + ) + R.output(out) + return out + + data = np.random.randn(1, 3, 8, 8).astype("float32") + patterns = [("tensorrt.image.resize2d", is_op("relax.image.resize2d")(wildcard(), wildcard()))] + _offload_and_compare(Cubic, {}, patterns, data) + + +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") + patterns = [("tensorrt.nn.silu", is_op("relax.nn.silu")(wildcard()))] + _offload_and_compare(Silu, {}, patterns, data) + + +def test_tensorrt_resize2d_partition_for_tensorrt(): + from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt + + @tvm.script.ir_module + class Model: + @R.function + def main( + data: R.Tensor((1, 8, 8, 8), "float32"), weight: R.Tensor((8, 8, 3, 3), "float32") + ): + with R.dataflow(): + conv = relax.op.nn.conv2d(data, weight, padding=1) + out = relax.op.image.resize2d( + conv, size=(16, 16), layout="NCHW", method="nearest_neighbor" + ) + R.output(out) + return out + + data = np.random.randn(1, 8, 8, 8).astype("float32") + weight = np.random.randn(8, 8, 3, 3).astype("float32") + ref = build_and_run(Model, [data, weight], "llvm", legalize=True) + + mod = relax.transform.BindParams("main", {"weight": weight})(Model) + mod = partition_for_tensorrt(mod) + codegen_fns = [ + fn + for fn in mod.functions.values() + if isinstance(fn, relax.Function) and fn.attrs is not None and "Codegen" in fn.attrs + ] + assert len(codegen_fns) == 1 + + mod = relax.transform.RunCodegen()(mod) + out = build_and_run(mod, [data], "cuda") + tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2) + + 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 c9981de37a69..0e8d403baed9 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 67dce2ba6531..793804a53374 100644 --- a/tests/python/relax/test_transform.py +++ b/tests/python/relax/test_transform.py @@ -272,6 +272,133 @@ def foo(x: R.Tensor(("m", "n"), "float32")): assert s2.op.global_symbol == "test.op.identity" +def test_call_dps_packed_rewrite_nested_tuple_output(): + """Nested outputs are flattened for the packed ABI and rebuilt afterwards.""" + 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) + before = builder.get() + + after = relax.transform.CallTIRRewrite()(before) + 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 len(packed_call.args) == 4 + assert packed_call.args[0].same_as(func.params[0]) + for arg, binding in zip(packed_call.args[1:], alloc_bindings): + assert arg.same_as(binding.var) + + rebuilt = block.bindings[4].value + assert isinstance(rebuilt, relax.Tuple) + assert rebuilt.fields[0].same_as(alloc_bindings[0].var) + assert isinstance(rebuilt.fields[1], relax.Tuple) + 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_dps_packed_rewrite_single_leaf_nested_tuple_output(): + """A single flattened output must still retain its nested tuple structure.""" + input_ty = relax.TensorType((2, 3), "float32") + output_leaf_ty = relax.TensorType((4,), "int32") + output_ty = tvm.ir.TupleType([tvm.ir.TupleType([output_leaf_ty])]) + + x = relax.Var("x", input_ty) + call = relax.Call( + tvm.ir.Op.get("relax.call_dps_packed"), + [relax.ExternFunc("test.op.single_nested_output"), 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) + before = builder.get() + + after = relax.transform.CallTIRRewrite()(before) + relax.analysis.well_formed(after) + func = after["main"] + block = func.body.blocks[0] + + alloc_binding = block.bindings[0] + assert alloc_binding.value.op.name == "relax.builtin.alloc_tensor" + tvm.ir.assert_structural_equal(alloc_binding.var.ty, output_leaf_ty) + + packed_call = block.bindings[1].value + assert packed_call.op.global_symbol == "test.op.single_nested_output" + assert len(packed_call.args) == 2 + assert packed_call.args[0].same_as(func.params[0]) + assert packed_call.args[1].same_as(alloc_binding.var) + + rebuilt = block.bindings[2].value + assert isinstance(rebuilt, relax.Tuple) + assert len(rebuilt.fields) == 1 + assert isinstance(rebuilt.fields[0], relax.Tuple) + assert len(rebuilt.fields[0].fields) == 1 + assert rebuilt.fields[0].fields[0].same_as(alloc_binding.var) + + +def test_call_dps_packed_rewrite_single_element_tuple_output(): + """A single-element tuple output must not be collapsed into its tensor field.""" + input_ty = relax.TensorType((2, 3), "float32") + output_leaf_ty = relax.TensorType((4,), "int32") + output_ty = tvm.ir.TupleType([output_leaf_ty]) + + x = relax.Var("x", input_ty) + call = relax.Call( + tvm.ir.Op.get("relax.call_dps_packed"), + [relax.ExternFunc("test.op.single_tuple_output"), 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) + before = builder.get() + + after = relax.transform.CallTIRRewrite()(before) + relax.analysis.well_formed(after) + func = after["main"] + block = func.body.blocks[0] + + alloc_binding = block.bindings[0] + assert alloc_binding.value.op.name == "relax.builtin.alloc_tensor" + tvm.ir.assert_structural_equal(alloc_binding.var.ty, output_leaf_ty) + + packed_call = block.bindings[1].value + assert packed_call.op.global_symbol == "test.op.single_tuple_output" + assert len(packed_call.args) == 2 + assert packed_call.args[0].same_as(func.params[0]) + assert packed_call.args[1].same_as(alloc_binding.var) + + rebuilt = block.bindings[2].value + assert isinstance(rebuilt, relax.Tuple) + assert len(rebuilt.fields) == 1 + assert rebuilt.fields[0].same_as(alloc_binding.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 7ef59cefbc4f..014833cf9638 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,262 @@ def main( check(mod, [("x.concat", pat_clip)], Expected2) +def _get_codegen_calls(mod): + """Return main-function codegen calls, keyed by backend name.""" + calls = {} + for block in mod["main"].body.blocks: + for binding in block.bindings: + value = binding.value + if not isinstance(value, relax.Call) or not isinstance(value.op, relax.GlobalVar): + continue + callee = mod[value.op] + if not isinstance(callee, relax.Function) or callee.attrs is None: + continue + if codegen := callee.attrs.get("Codegen"): + calls[str(codegen)] = (binding.var, value) + return calls + + +def _find_tuple_get_item(mod, tuple_value, index): + """Find the variable bound to a specific TupleGetItem in main.""" + matches = [] + for block in mod["main"].body.blocks: + for binding in block.bindings: + value = binding.value + if ( + isinstance(value, relax.TupleGetItem) + and value.index == index + and value.tuple_value.same_as(tuple_value) + ): + matches.append(binding.var) + assert len(matches) == 1 + return matches[0] + + +def test_tuple_output_remap_is_scoped_to_producer_group(): + """One tuple-producing group must not suppress another group's output remapping.""" + + @I.ir_module + class Before: + @R.function(private=True) + def relu( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tensor((2, 4), "float32"): + R.func_attr({"Composite": "compiler_A.relu", "Primitive": True}) + return R.nn.relu(x) + + @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_B.split", "Primitive": True}) + return R.split(x, indices_or_sections=2, axis=0) + + @R.function + def main( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tuple( + R.Tensor((2, 4), "float32"), + R.Tensor((2, 4), "float32"), + R.Tensor((1, 4), "float32"), + ): + cls = Before + with R.dataflow(): + first = cls.relu(x) + second = cls.relu(first) + parts = cls.split(x) + left = parts[0] + out = (first, second, left) + R.output(out) + return out + + after = relax.transform.MergeCompositeFunctions()(Before) + relax.analysis.well_formed(after) + assert not relax.analysis.free_vars(after["main"]) + assert set(_get_codegen_calls(after)) == {"compiler_A", "compiler_B"} + + +def test_multiple_boundary_outputs_preserve_nested_tuple_indices(): + """Fill missing output remaps without replacing existing nested tuple indices.""" + + @I.ir_module + class Before: + @R.function(private=True) + def relu( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tensor((2, 4), "float32"): + R.func_attr({"Composite": "compiler_A.relu", "Primitive": True}) + return R.nn.relu(x) + + @R.function(private=True) + def split_a( + 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 split_b( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tuple(R.Tensor((1, 4), "float32"), R.Tensor((1, 4), "float32")): + R.func_attr({"Composite": "compiler_B.split", "Primitive": True}) + return R.split(x, indices_or_sections=2, axis=0) + + @R.function + def main( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tuple(R.Tensor((1, 4), "float32"), R.Tensor((1, 4), "float32")): + cls = Before + with R.dataflow(): + producer = cls.relu(x) + parts_a = cls.split_a(producer) + side = cls.relu(producer) + a_right = parts_a[1] + parts_b = cls.split_b(side) + b_left = parts_b[0] + out = (a_right, b_left) + R.output(out) + return out + + after = relax.transform.MergeCompositeFunctions()(Before) + relax.analysis.well_formed(after) + assert not relax.analysis.free_vars(after["main"]) + + calls = _get_codegen_calls(after) + result_a, _ = calls["compiler_A"] + result_b, call_b = calls["compiler_B"] + + parts_a = _find_tuple_get_item(after, result_a, 0) + side = _find_tuple_get_item(after, result_a, 1) + _find_tuple_get_item(after, parts_a, 1) + _find_tuple_get_item(after, result_b, 0) + assert call_b.args[0].same_as(side) + + +def test_single_nested_tuple_boundary_output_is_not_unwrapped(): + """A sole nested-tuple boundary output maps to the call, not call_result[0].""" + + @I.ir_module + class Before: + @R.function(private=True) + def make_nested( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tuple( + R.Tuple(R.Tensor((1, 4), "float32"), R.Tensor((1, 4), "float32")), + R.Tensor((2, 4), "float32"), + ): + R.func_attr({"Composite": "compiler_A.make_nested", "Primitive": True}) + parts = R.split(x, indices_or_sections=2, axis=0) + return (parts, x) + + @R.function + def main( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tensor((1, 4), "float32"): + cls = Before + with R.dataflow(): + nested = cls.make_nested(x) + pair = nested[0] + right = pair[1] + R.output(right) + return right + + after = relax.transform.MergeCompositeFunctions()(Before) + relax.analysis.well_formed(after) + assert not relax.analysis.free_vars(after["main"]) + + result, _ = _get_codegen_calls(after)["compiler_A"] + pair = _find_tuple_get_item(after, result, 0) + right = _find_tuple_get_item(after, pair, 1) + assert after["main"].body.body.same_as(right) + + +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"]) + + grouped_result, _ = _get_codegen_calls(after)["compiler_A"] + assert not isinstance(grouped_result, relax.DataflowVar) + assert after["main"].body.body.same_as(grouped_result) + + +def test_multiple_boundary_outputs_preserve_dataflow_status(): + """Only boundary outputs that escape the dataflow block become output variables.""" + + @I.ir_module + class Before: + @R.function(private=True) + def relu( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tensor((2, 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.Tensor((2, 4), "float32"), R.Tensor((2, 4), "float32")): + cls = Before + with R.dataflow(): + first = cls.relu(x) + out = cls.relu(first) + side = R.sigmoid(first) + R.output(out, side) + return (out, side) + + after = relax.transform.MergeCompositeFunctions()(Before) + relax.analysis.well_formed(after) + assert not relax.analysis.free_vars(after["main"]) + + grouped_result, _ = _get_codegen_calls(after)["compiler_A"] + assert isinstance(grouped_result, relax.DataflowVar) + + tuple_get_bindings = [ + binding + for block in after["main"].body.blocks + for binding in block.bindings + if isinstance(binding.value, relax.TupleGetItem) + and binding.value.tuple_value.same_as(grouped_result) + ] + assert len(tuple_get_bindings) == 2 + + output_var = after["main"].body.body.fields[0] + assert not isinstance(output_var, relax.DataflowVar) + assert any(binding.var.same_as(output_var) for binding in tuple_get_bindings) + + sigmoid_calls = [ + binding.value + for block in after["main"].body.blocks + for binding in block.bindings + if isinstance(binding.value, relax.Call) + and isinstance(binding.value.op, tvm.ir.Op) + and binding.value.op.name == "relax.sigmoid" + ] + assert len(sigmoid_calls) == 1 + internal_var = sigmoid_calls[0].args[0] + assert isinstance(internal_var, relax.DataflowVar) + assert any(binding.var.same_as(internal_var) for binding in tuple_get_bindings) + + 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 b1ce09799a64..7f4f051e62c6 100644 --- a/tests/python/relax/test_transform_merge_composite_functions.py +++ b/tests/python/relax/test_transform_merge_composite_functions.py @@ -1221,5 +1221,66 @@ def composite_lambda( tvm.ir.assert_structural_equal(Expected, After) +def test_tuple_get_item_dependency_prevents_cyclic_merge(): + """TupleGetItem dependencies must be considered when merging composite regions.""" + + @tvm.script.ir_module + class Before: + @R.function(private=True) + def relu( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tensor((2, 4), "float32"): + R.func_attr({"Composite": "compiler_A.relu", "Primitive": True}) + return R.nn.relu(x) + + @R.function(private=True) + def relu_half( + 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(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 concat( + x: R.Tensor((2, 4), "float32"), + y: R.Tensor((1, 4), "float32"), + ) -> R.Tensor((3, 4), "float32"): + R.func_attr({"Composite": "compiler_A.concat", "Primitive": True}) + return R.concat((x, y), axis=0) + + @R.function + def main( + x: R.Tensor((2, 4), "float32"), + ) -> R.Tensor((3, 4), "float32"): + cls = Before + with R.dataflow(): + producer = cls.relu(x) + parts = cls.split(producer) + right = parts[1] + side = cls.relu(producer) + right_out = cls.relu_half(right) + out = cls.concat(side, right_out) + R.output(out) + return out + + after = relax.transform.MergeCompositeFunctions()(Before) + codegen_regions = [ + 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" + ] + assert len(codegen_regions) == 2 + relax.analysis.well_formed(after) + + if __name__ == "__main__": pytest.main([__file__])