[BUG] Segmentation fault during executing a module that uses pass tir.transform.DecorateDeviceScope

[syang-ng]

Hi there, we are trying to apply some tir passes for our program. We encountered the problem on pass tir.transform.DecorateDeviceScope:

The code to reproduce the problem is as below:

# demo.py
import tvm
import numpy as np
from tvm import tir, te

n = te.size_var("n")
m = te.size_var("m")
A = te.placeholder((n, n), name="A", dtype="int32")

T = te.compute((m, m), lambda i, j: A[i][j])
s = te.create_schedule(T.op)
ir_m = tvm.lower(s, [A, T])

inputs = [tvm.nd.array(np.random.uniform(0, 100, size=(32, 32)).astype("int32"))]
output = tvm.nd.empty((32, 32), "int32")

with tvm.transform.PassContext(opt_level=4):
    opt = tvm.transform.Sequential(
        [tir.transform.DecorateDeviceScope()]
    )
    mod = opt(ir_m)

    opt_execute = tvm.build(mod, [*inputs, output], tvm.target.Target("llvm"))

    # segmentation fault here
    opt_execute(inputs[0], output)

A segmentation fault occurs during this program execution.

(tvm-build) ➜  ~ python3 demo.py
> INIT:: # BB in this DSO: 588343; # BB total: 588343
[1]    25204 segmentation fault (core dumped)  python3 demo.py

It seems that after applying the pass tir.transform.DecorateDeviceScope(), a bug is left in the final built module opt_execute and then the error occurs during its execution.

Environment Information

OS/Compiler

lsb_release -a
No LSB modules are available.
Distributor ID: Ubuntu
Description:    Ubuntu 18.04.5 LTS
Release:        18.04
Codename:       bionic

clang++-12 --version
Ubuntu clang version 12.0.1-++20210626052611+b7c7b42db1d1-1~exp1~20210626033351.133
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin

Config

In addition to default one, the following options are set as 'ON':

• USE_LLVM

Linked libraries

ldd libtvm.so
        linux-vdso.so.1 (0x00007ffd9bbb0000)
        libLLVM-12.so.1 => /usr/lib/x86_64-linux-gnu/libLLVM-12.so.1 (0x00007fd7ee4fa000)
        libdl.so.2 => /lib/x86_64-linux-gnu/libdl.so.2 (0x00007fd7ee2f6000)
        libmemcov.so => /home/scam/tvm-cov/build/_deps/memcov-build/libmemcov.so (0x00007fd7ee0f4000)
        libstdc++.so.6 => /usr/lib/x86_64-linux-gnu/libstdc++.so.6 (0x00007fd7edd6b000)
        libm.so.6 => /lib/x86_64-linux-gnu/libm.so.6 (0x00007fd7ed9cd000)
        libgcc_s.so.1 => /lib/x86_64-linux-gnu/libgcc_s.so.1 (0x00007fd7ed7b5000)
        libpthread.so.0 => /lib/x86_64-linux-gnu/libpthread.so.0 (0x00007fd7ed596000)
        libc.so.6 => /lib/x86_64-linux-gnu/libc.so.6 (0x00007fd7ed1a5000)
        /lib64/ld-linux-x86-64.so.2 (0x00007fd7f88ca000)
        libffi.so.6 => /usr/lib/x86_64-linux-gnu/libffi.so.6 (0x00007fd7ecf9d000)
        libedit.so.2 => /usr/lib/x86_64-linux-gnu/libedit.so.2 (0x00007fd7ecd66000)
        librt.so.1 => /lib/x86_64-linux-gnu/librt.so.1 (0x00007fd7ecb5e000)
        libz.so.1 => /lib/x86_64-linux-gnu/libz.so.1 (0x00007fd7ec941000)
        libtinfo.so.5 => /lib/x86_64-linux-gnu/libtinfo.so.5 (0x00007fd7ec717000)

[Lunderberg]

Confirmed that the issue is still present in main, reproducible with the test script. Segfault occurs during execution of the generated kernel.