[Relax][Frontend][ONNX] Support dynamic Range bounds

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -3724,24 +3724,76 @@ def _impl_v12(cls, bb, inputs, attr, params):
         start = get_constant(inputs[0], params)
         limit = get_constant(inputs[1], params)
         delta = get_constant(inputs[2], params)
-        out_dtype = start.ty.dtype
 
-        if isinstance(start, relax.Constant):
-            start = start.data.numpy().tolist()
+        def get_scalar_dtype(x):
+            if tvm.ir.is_prim_expr(x):
+                return str(getattr(x, "dtype", None) or x.ty)
+            return str(x.ty.dtype)
 
-        if isinstance(limit, relax.Constant):
-            limit = limit.data.numpy().tolist()
+        out_dtype = get_scalar_dtype(start)
 
-        assert isinstance(delta, relax.Constant), "Constant delta required for Range."
-        step = delta.data.numpy().tolist()
+        def get_scalar_value(x):
+            if isinstance(x, relax.Constant):
+                value = x.data.numpy()
+                if value.size != 1:
+                    raise ValueError("Range scalar input must have exactly one element.")
+                return value.item()
+            return x
+
+        start = get_scalar_value(start)
+        limit = get_scalar_value(limit)
+        delta = get_scalar_value(delta)
+
+        def is_dynamic_scalar(x):
+            return tvm.ir.is_prim_expr(x) or isinstance(x, relax.Expr)
 
-        # If all inputs are constant, compute directly.
-        if isinstance(start, int) and isinstance(limit, int):
-            out_range = _np.arange(start=start, stop=limit, step=step)
+        if not any(is_dynamic_scalar(x) for x in [start, limit, delta]):
+            out_range = _np.arange(start=start, stop=limit, step=delta)
             return relax.const(out_range, out_dtype)
 
-        # Otherwise compute in graph.
-        return relax.op.arange(start, limit, step, out_dtype)
+        out_dtype_is_float = _relax_dtype_is_floating_point(out_dtype)
+        count_dtype = "float64" if out_dtype_is_float else "int64"
+
+        def scalar_expr(x, dtype):
+            if tvm.ir.is_prim_expr(x):
+                expr_dtype = str(getattr(x, "dtype", None) or x.ty)
+                if expr_dtype != "int64":
+                    x = tirx.Cast("int64", x)
+                x = bb.normalize(relax.op.shape_to_tensor(relax.ShapeExpr([x])))
+                x = bb.normalize(relax.op.reshape(x, ()))
+                if dtype != "int64":
+                    x = bb.normalize(relax.op.astype(x, dtype))
+                return x
+            if isinstance(x, relax.Expr):
+                if str(x.ty.dtype) == dtype:
+                    return x
+                return bb.normalize(relax.op.astype(x, dtype))
+            return relax.const(x, dtype)
+
+        start_count = scalar_expr(start, count_dtype)
+        limit_count = scalar_expr(limit, count_dtype)
+        delta_count = scalar_expr(delta, count_dtype)
+
+        if out_dtype_is_float:
+            count = relax.op.ceil(
+                relax.op.divide(relax.op.subtract(limit_count, start_count), delta_count)
+            )
+        else:
+            count = relax.op.negative(
+                relax.op.floor_divide(relax.op.subtract(start_count, limit_count), delta_count)
+            )
+
+        count = bb.normalize(relax.op.maximum(count, relax.const(0, count_dtype)))
+        count = bb.normalize(relax.op.astype(count, "int64"))
+        count = bb.normalize(relax.op.reshape(count, (1,)))
+        range_len = _tensor_to_shape_expr(bb, count, 1, "range_len").values[0]
+
+        positions = bb.normalize(
+            relax.op.astype(relax.op.arange(0, range_len, 1, "int64"), out_dtype)
+        )
+        start_value = scalar_expr(start, out_dtype)
+        delta_value = scalar_expr(delta, out_dtype)
+        return relax.op.add(relax.op.multiply(positions, delta_value), start_value)
 
 
 class InstanceNormalization(OnnxOpConverter):

tests/python/relax/test_frontend_onnx.py

@@ -8216,6 +8216,141 @@ def main(
     tvm.ir.assert_structural_equal(tvm_model, Expected)
 
 
+@pytest.mark.parametrize(
+    "start, limit, delta, tensor_dtype, np_dtype",
+    [
+        (0, 6, 2, TensorProto.INT64, np.int64),
+        (8, 0, -2, TensorProto.INT64, np.int64),
+        (5, 1, 1, TensorProto.INT64, np.int64),
+        (0, 7, 2, TensorProto.INT32, np.int32),
+        (0.0, 1.0, 0.25, TensorProto.FLOAT, np.float32),
+        (1.0, -1.0, -0.5, TensorProto.FLOAT, np.float32),
+        (0.0, 0.3, 0.1, TensorProto.FLOAT, np.float32),
+    ],
+)
+def test_range_dynamic_scalar_inputs(start, limit, delta, tensor_dtype, np_dtype):
+    range_node = helper.make_node(
+        "Range",
+        ["start", "limit", "delta"],
+        ["output"],
+    )
+
+    graph = helper.make_graph(
+        [range_node],
+        "range_dynamic_scalar_inputs_test",
+        inputs=[
+            helper.make_tensor_value_info("start", tensor_dtype, []),
+            helper.make_tensor_value_info("limit", tensor_dtype, []),
+            helper.make_tensor_value_info("delta", tensor_dtype, []),
+        ],
+        outputs=[
+            helper.make_tensor_value_info("output", tensor_dtype, ["range_len"]),
+        ],
+    )
+
+    model = helper.make_model(graph, producer_name="range_dynamic_scalar_inputs_test")
+    check_correctness(
+        model,
+        inputs={
+            "start": np.array(start, dtype=np_dtype),
+            "limit": np.array(limit, dtype=np_dtype),
+            "delta": np.array(delta, dtype=np_dtype),
+        },
+        opset=12,
+        check_dtypes=True,
+    )
+
+
+def test_range_mixed_tensor_and_primexpr_limit():
+    shape = helper.make_node("Shape", ["x"], ["x_shape"])
+    axis = make_constant_node("axis", TensorProto.INT64, [], [1])
+    gather = helper.make_node("Gather", ["x_shape", "axis"], ["limit_int"])
+    cast = helper.make_node("Cast", ["limit_int"], ["limit"], to=TensorProto.FLOAT)
+    delta = make_constant_node("delta", TensorProto.FLOAT, [], [1.0])
+    range_node = helper.make_node(
+        "Range",
+        ["start", "limit", "delta"],
+        ["output"],
+    )
+
+    graph = helper.make_graph(
+        [shape, axis, gather, cast, delta, range_node],
+        "range_mixed_tensor_and_primexpr_limit_test",
+        inputs=[
+            helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, "range_len"]),
+            helper.make_tensor_value_info("start", TensorProto.FLOAT, []),
+        ],
+        outputs=[
+            helper.make_tensor_value_info("output", TensorProto.FLOAT, ["range_len"]),
+        ],
+    )
+
+    model = helper.make_model(
+        graph,
+        producer_name="range_mixed_tensor_and_primexpr_limit_test",
+        opset_imports=[helper.make_opsetid("", 17)],
+    )
+    model.ir_version = 8
+    check_correctness(
+        model,
+        inputs={
+            "x": np.ones((1, 4), dtype=np.float32),
+            "start": np.array(0.0, dtype=np.float32),
+        },
+        opset=17,
+        check_dtypes=True,
+    )
+
+
+@pytest.mark.parametrize(
+    "start_from_dim, limit_from_dim, delta",
+    [
+        (True, False, -1),
+        (False, True, -1),
+    ],
+)
+def test_range_primexpr_negative_and_empty(start_from_dim, limit_from_dim, delta):
+    shape = helper.make_node("Shape", ["x"], ["x_shape"])
+    axis = make_constant_node("axis", TensorProto.INT64, [], [1])
+    gather = helper.make_node("Gather", ["x_shape", "axis"], ["dim"])
+    start = make_constant_node("start", TensorProto.INT64, [], [0])
+    limit = make_constant_node("limit", TensorProto.INT64, [], [0])
+    delta_node = make_constant_node("delta", TensorProto.INT64, [], [delta])
+
+    range_inputs = [
+        "dim" if start_from_dim else "start",
+        "dim" if limit_from_dim else "limit",
+        "delta",
+    ]
+    range_node = helper.make_node("Range", range_inputs, ["output"])
+
+    graph = helper.make_graph(
+        [shape, axis, gather, start, limit, delta_node, range_node],
+        "range_primexpr_negative_and_empty_test",
+        inputs=[
+            helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, "range_len"]),
+        ],
+        outputs=[
+            helper.make_tensor_value_info("output", TensorProto.INT64, ["output_len"]),
+        ],
+    )
+
+    model = helper.make_model(
+        graph,
+        producer_name="range_primexpr_negative_and_empty_test",
+        opset_imports=[helper.make_opsetid("", 17)],
+    )
+    model.ir_version = 8
+    check_correctness(
+        model,
+        inputs={
+            "x": np.ones((1, 4), dtype=np.float32),
+        },
+        opset=17,
+        check_dtypes=True,
+    )
+
+
 def test_batch_norm():
     batch_norm_node = helper.make_node(
         "BatchNormalization", ["x", "s", "bias", "mean", "var"], ["y"], epsilon=1e-2