[Relax][TensorRT] Fix YOLO BYOC offload and partitioning gaps

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"),

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):

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

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

src/relax/transform/merge_composite_functions.cc

@@ -89,7 +89,6 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator<Group*> {
       // 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>();
       }
@@ -166,6 +165,18 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator<Group*> {
     return group;
   }
 
+  Group* VisitExpr_(const TupleNode* tuple) {
+    Group* group = arena_->make<Group>();
+    UpdateGroupDependencies(group, tuple->fields);
+    return group;
+  }
+
+  Group* VisitExpr_(const TupleGetItemNode* tuple_get_item) {
+    Group* group = arena_->make<Group>();
+    UpdateGroupDependencies(group, {tuple_get_item->tuple});
+    return group;
+  }
+
  private:
   ffi::Optional<ffi::String> GetCodegenName(const Expr& callee) {
     auto const* gvar = callee.as<GlobalVarNode>();

src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc

@@ -27,8 +27,10 @@
 #include <tvm/runtime/logging.h>
 #include <tvm/runtime/tensor.h>
 
+#include <algorithm>
 #include <memory>
 #include <string>
+#include <vector>
 
 #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<std::string, std::unique_ptr<TensorRTOpConverter>>& map) {
+  std::vector<std::string> 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<std::string, std::unique_ptr<TensorRTOpConverter>>& 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())

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<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 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<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 (method == "nearest_neighbor") {
+      static const std::unordered_map<std::string, nvinfer1::ResizeRoundMode> 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<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 +1382,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 +1448,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));

tests/python/relax/test_codegen_tensorrt.py

@@ -640,5 +640,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()

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():