[Bug] Flaky Behavior in Relax VM Execution: "ValueError: negative dimensions are not allowed" and "Floating point exception (core dumped)"

[Thrsu]

When running a Relax program that involves dynamic strided slicing with specific input values, I encounter three types of behavior:

1. The code runs without errors.
2. The error ValueError: negative dimensions are not allowed is raised.
3. A Floating point exception (core dumped) error causes the program to crash.

This inconsistent behavior appears to depend on the values of the input tensors for begin, end, and strides. It seems that the shape function or dynamic strided slicing logic isn't handling all input cases properly, leading to negative dimensions or floating-point exceptions.

Expected behavior

The program should consistently run without errors, handling the input values appropriately for all valid cases.

Actual behavior

1. Run well
2. ValueError: negative dimensions are not allowed
3. Floating point exception (core dumped)

Steps to reproduce

import tvm
from tvm import relax
import numpy as np
import os
from tvm.script import ir as I
from tvm.script import tir as T
from tvm.script import relax as R

@I.ir_module
class Module:
    @T.prim_func(private=True)
    def dynamic_strided_slice(var_x: T.handle, begin: T.Buffer((T.int64(2),), "int64"), end: T.Buffer((T.int64(2),), "int64"), strides: T.Buffer((T.int64(2),), "int64"), var_T_strided_slice_dynamic: T.handle):
        T.func_attr({"tir.noalias": T.bool(True)})
        n = T.int64()
        x = T.match_buffer(var_x, (T.int64(10), n))
        s, s_1 = T.int64(), T.int64()
        T_strided_slice_dynamic = T.match_buffer(var_T_strided_slice_dynamic, (s, s_1))
        # with T.block("root"):
        for ax0, ax1 in T.grid(s, s_1):
            with T.block("T_strided_slice_dynamic"):
                v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                T.reads(x[T.min(begin[T.int64(0)], T.int64(9)) + v_ax0 * strides[T.int64(0)], T.min(begin[T.int64(1)], n - T.int64(1)) + v_ax1 * strides[T.int64(1)]], begin[T.int64(0):T.int64(2)], strides[T.int64(0):T.int64(2)])
                T.writes(T_strided_slice_dynamic[v_ax0, v_ax1])
                T_strided_slice_dynamic[v_ax0, v_ax1] = x[T.min(begin[T.int64(0)], T.int64(9)) + v_ax0 * strides[T.int64(0)], T.min(begin[T.int64(1)], n - T.int64(1)) + v_ax1 * strides[T.int64(1)]]

    @T.prim_func(private=True)
    def shape_func(var_x: T.handle, begin: T.Buffer((T.int64(2),), "int64"), end: T.Buffer((T.int64(2),), "int64"), strides: T.Buffer((T.int64(2),), "int64"), T_shape_func_strided_slice_dynamic: T.Buffer((T.int64(2),), "int64")):
        T.func_attr({"tir.noalias": T.bool(True)})
        n = T.int64()
        x = T.match_buffer(var_x, (T.int64(10), n))
        # with T.block("root"):
        for i in range(T.int64(2)):
            with T.block("T_shape_func_strided_slice_dynamic"):
                v_i = T.axis.spatial(T.int64(2), i)
                T.reads(strides[v_i], begin[v_i], end[v_i])
                T.writes(T_shape_func_strided_slice_dynamic[v_i])
                T_shape_func_strided_slice_dynamic[v_i] = T.Select(strides[v_i] < T.int64(0), (T.min(T.max(T.Select(begin[v_i] < T.int64(0), begin[v_i] + T.Select(v_i == T.int64(1), n, T.Select(v_i == T.int64(0), T.int64(10), T.int64(-1))), begin[v_i]), T.int64(-1)), T.Select(v_i == T.int64(1), n, T.Select(v_i == T.int64(0), T.int64(10), T.int64(-1))) - T.int64(1)) - T.min(T.max(T.Select(end[v_i] < T.int64(0), end[v_i] + T.Select(v_i == T.int64(1), n, T.Select(v_i == T.int64(0), T.int64(10), T.int64(-1))), end[v_i]), T.int64(-1)), T.Select(v_i == T.int64(1), n, T.Select(v_i == T.int64(0), T.int64(10), T.int64(-1))) - T.int64(1)) - strides[v_i] - T.int64(1)) // (strides[v_i] * T.int64(-1)), (T.min(T.max(T.Select(end[v_i] < T.int64(0), end[v_i] + T.Select(v_i == T.int64(1), n, T.Select(v_i == T.int64(0), T.int64(10), T.int64(-1))), end[v_i]), T.int64(0)), T.Select(v_i == T.int64(1), n, T.Select(v_i == T.int64(0), T.int64(10), T.int64(-1)))) + strides[v_i] - T.min(T.max(T.Select(begin[v_i] < T.int64(0), begin[v_i] + T.Select(v_i == T.int64(1), n, T.Select(v_i == T.int64(0), T.int64(10), T.int64(-1))), begin[v_i]), T.int64(0)), T.Select(v_i == T.int64(1), n, T.Select(v_i == T.int64(0), T.int64(10), T.int64(-1)))) - T.int64(1)) // strides[v_i])

    @R.function
    def main(x: R.Tensor((10, "n"), dtype="float32"), begin: R.Tensor((2,), dtype="int64"), end: R.Tensor((2,), dtype="int64"), strides: R.Tensor((2,), dtype="int64")) -> R.Tensor(dtype="float32", ndim=2):
        n = T.int64()
        cls = Module
        s = T.int64()
        s_1 = T.int64()
        gv = R.call_tir(cls.shape_func, (x, begin, end, strides), out_sinfo=R.Tensor((2,), dtype="int64"))
        gv1: R.Shape(ndim=2) = R.call_pure_packed("vm.builtin.tensor_to_shape", gv, sinfo_args=(R.Shape(ndim=2),))
        gv2: R.Shape([s, s_1]) = R.match_cast(gv1, R.Shape([s, s_1]))
        gv_1 = R.call_tir(cls.dynamic_strided_slice, (x, begin, end, strides), out_sinfo=R.Tensor((s, s_1), dtype="float32"))
        return gv_1

mod = Module

def compile_mod(mod, func_name, target, *inputs):
    with tvm.transform.PassContext(opt_level=4): 
        ex = relax.build(mod, target='llvm')
    vm = relax.VirtualMachine(ex, tvm.cpu())
    mod_outputs = vm[f'{func_name}'](*inputs)
    print(mod_outputs)

input_0 = tvm.nd.array(10 * np.random.random([10, 1]).astype('float32'))
input_1 = tvm.nd.array(np.random.randint(10, size=[2]).astype('int64'))
input_2 = tvm.nd.array(np.random.randint(10, size=[2]).astype('int64'))
input_3 = tvm.nd.array(np.random.randint(10, size=[2]).astype('int64'))
compile_mod(mod, 'main', 'llvm', input_0,input_1,input_2,input_3)

CC @Lunderberg @tqchen @junrushao