search

facebookincubator / AITemplate

AITemplate is a Python framework which renders neural network into high performance CUDA/HIP C++ code. Specialized for FP16 TensorCore (NVIDIA GPU) and MatrixCore (AMD GPU) inference.
Apache License 2.0
4.55k stars 369 forks source link

The input/output of each kernel in stable diffustion #835

Open chsungen opened 1 year ago

chsungen commented 1 year ago

Hello,

Thanks for the very usable repo! Now I am trying to get the input/output shape of each kernel. For stable diffusion, I already found that after compile, I can get the shape in /AITemplate/examples/05_stable_diffusion/scripts/tmp/AutoencoderKL/model-generated.h (take vae model as an example). However, some of the information is still missing. For example,

image

Here the input shape still show some dimensions like "upsampling2d_79_0_dim_1", "upsampling2d_79_0_dim_2", ..., etc. And when I take a look for this variables, I found in the same file, it shows 0, please see the screen shot below.

image image

Is there any way to get these information? Thanks!

ipiszy commented 1 year ago

Hi @chsungen , you could set env "LOGLEVEL=DEBUG" when you use AIT to compile. This will generate a lot of debug files for each optimization step. e.g. "memory_planning_pseudo_code.txt" and "memory_planning_graph.txt" have information about input / output tensor shapes. "memory_planning_graph_vis.html" is a graph visualization.