diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index b43969e86639..d0025afeff60 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -408,6 +408,9 @@ bool HasReshapePattern(const PrimFunc& func) { ffi::Map var_range; for (const IterVar& v : block->iter_vars) { + if (ana_->CanProveEqual(v->dom->extent, IntImm::Int64(/*value=*/0))) { + return; + } ana_->Bind(v->var, Range::FromMinExtent(v->dom->min, v->dom->extent)); var_range.Set(v->var, Range::FromMinExtent(v->dom->min, v->dom->extent)); } diff --git a/tests/python/relax/test_analysis.py b/tests/python/relax/test_analysis.py index 7be99442aced..ecf5675109dd 100644 --- a/tests/python/relax/test_analysis.py +++ b/tests/python/relax/test_analysis.py @@ -811,5 +811,31 @@ def reduction(A: T.Buffer((4, 4), "float32"), B: T.Buffer((4,), "float32")): assert not has_reshape_pattern(reduction) +def test_reshape_pattern_reject_zero_extent_slice(): + @T.prim_func(s_tir=True) + def strided_slice(A: T.Buffer((3,), "float32"), T_out: T.Buffer((0,), "float32")): + for i in T.serial(0): + with T.sblock("T_strided_slice"): + vi = T.axis.spatial(0, i) + T.reads(A[vi]) + T.writes(T_out[vi]) + T_out[vi] = A[vi] + + assert not has_reshape_pattern(strided_slice) + + +def test_reshape_pattern_reject_zero_extent_multi_dim(): + @T.prim_func(s_tir=True) + def reshape_empty(A: T.Buffer((0, 4), "float32"), T_out: T.Buffer((0, 2, 2), "float32")): + for i0, i1, i2 in T.grid(0, 2, 2): + with T.sblock("T_reshape"): + v0, v1, v2 = T.axis.remap("SSS", [i0, i1, i2]) + T.reads(A[v0, v1 * 2 + v2]) + T.writes(T_out[v0, v1, v2]) + T_out[v0, v1, v2] = A[v0, v1 * 2 + v2] + + assert not has_reshape_pattern(reshape_empty) + + if __name__ == "__main__": tvm.testing.main()