Inquiry About Obtaining PTX Assembly Code from TVM

Questions
#1

Hello,

I’m currently studying how neural network models are executed on GPUs. As part of my studies, I am particularly interested in obtaining the PTX assembly code for network models after they have been compiled using TVM.

I have been exploring the capabilities of TVM and have successfully compiled neural network models to run on GPUs.

I’d like to know how to access the PTX assembly code generated after compiling network models using TVM. Could you please guide me on how to obtain this PTX code?

#2

As I know about tvm, for single op, tvm can compile high level ir to cuda c source code, and then use the nvcc to compile the source into an executable. To extract the PTX code, you can retrieve the source code and use appropriate tools to compile it into PTX.

cuda_mod = tvm.build(sch.mod, target="cuda")
c_source_code = cuda_mod.imported_modules[0].get_source()

# Save the CUDA C source code to a file  
with open("cuda_source_code.cu", "w") as output_file:  
    output_file.write(c_source_code)  
# nvcc -ptx -o output.ptx cuda_source_code.cu

Btw, I think the simplest way to obtain the ptx code is to use Nvidia Nsight Compute to profile a TVM CUDA runtime. You can get the ptx code in the source window.

Regarding the network, I don’t see any built-in pass that directly generates a PTX source file.

TVM unroll is different from real code
#3

Another way is to override tvm_callback_cuda_compile

def my_cuda_compile(code, target):
     # Code to print or write to file
     return nvcc.compile_cuda(code)

tvm.register_func("tvm_callback_cuda_compile", my_cuda_compile, override=True)
#4

hi all,

Is that LLVM NVPTX backend is not supported any more, or it’s not an optimal codegen path working together with tir and metaschedule compared to cuda-c?