Segfault when printing mod in a prim_func_pass with UMA

I am playing with the UMA tutorial and trying to get a better feel for it, I have added a MyAcceleratorDensePass prim_func_pass that just prints the func and mod and returns the func:

@tvm.tir.transform.prim_func_pass(opt_level=2)
class MyAcceleratorDensePass:
    _EXTERNAL_FUNCTION_NAME = "my_accelerator_dense"
    _TVM_BLOCK_MATCH_NAME = "dense"

    def transform_function(
        self, func: tvm.tir.PrimFunc, mod: tvm.ir.IRModule, ctx: tvm.ir.transform.PassContext
    ) -> tvm.tir.PrimFunc:
        print(80*'-')
        print('MyAcceleratorDensePass')
        print(mod)
        print(80*'-')
        print(func)
        print(80*'-')
        return func

Which I register form my UMABackend using self._register_tir_pass(PassPhase.TIR_PHASE_0, MyAcceleratorDensePass()).

I get a segfault when I print(mod), if I remove it then no segfault.

Is it a bug or did I do something wrong?

Hi, I had this problem too, but I tried printing the “func” to find the required info which I needed.

1 Like