[Bug] Unexpected Segmentation Fault with Relax R.astype When Tensor Shape is Miswritten

[Thrsu]

There seems to be an issue when using R.astype in TVM's Relax script with a tensor shape that does not include a comma in the shape definition. Specifically, when a tensor is defined with the shape (1) instead of (1,), a segmentation fault occurs. The script runs without issue when the shape is correctly specified as (1,).

Expected behavior

A segmentation fault should not occur in this scenario.

Actual behavior

error: object of type 'int' has no len()
 --> seg_bug.py:14:17
    |  
 14 |              gv: R.Tensor((1), dtype="float32") = R.astype(lv, dtype="float32")  # segfault
    |                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                            
Segmentation fault (core dumped)

Steps to reproduce

import tvm
from tvm.script import ir as I
from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def main(z: R.Tensor((1, 4, 64, 64), dtype="float32")) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
        with R.dataflow():
            lv: R.Tensor((1, 4, 64, 64), dtype="float16") = R.astype(z, dtype="float16")
            # gv: R.Tensor((1,), dtype="float32") = R.astype(lv, dtype="float32")  # This works fine
            gv: R.Tensor((1), dtype="float32") = R.astype(lv, dtype="float32")  # Segfault here
            R.output(gv)
        return gv

mod = Module
mod.show()
mod = tvm.relax.transform.LegalizeOps()(mod)
mod.show()

cc @Lunderberg @tqchen @junrushao

[Lunderberg]

Looks like this is a bug in the TVMScript parsing. If an error occurred, it prints using the diagnostic renderer, and then overwrites the DiagnosticRenderer with a default-constructed instance. This default-constructed instance has a null pointer for its callback function, so a second call to DiagnosticContext::Render triggers the segfault. This looks like it is intended to prevent double-printing of error messages, but double-printing error messages is much worse than a segfault.

But we also shouldn't be calling it multiple times in the first place. There's a common pattern in the parser where it will wrap exceptions in DiagnosticError, but re-raise exceptions that are already a DiagnosticError. This requires every such location to include except DiagnosticError: raise, and can easily be missed.

Can you retry with #17373 applied? It should resolve both sources of the bug.

[Thrsu]

I applied this PR and the issue was resolved. I really appreciate your thorough explanation of the bug and the effort you put into fixing it. Thank you again for your help!