ControlNet-trt优化总结3：使用multi-stream和cuda-graph构建并行流水线

上节谈到使用TRT-API来构建网络，在这一节中总结一些trick来提升模型的运行效率，这些trick在所有的trt优化中均可使用，主要有以下几点：

使用cuda_graph减少kernel间的启动间隙
使用Mutil-stream增加异步

cuda_graph

cuda_graph的引入是为了解决kernel间launch的间隙时间问题的，尤其是有一堆小kernel,每个kernel启动也会带来一些开销，如果这些kernel足够多，那么就可能会影响系统的整体性能，cuda_graph的引入就是为了解决这个问题的，它会将stream内的kernel视为一整个graph，从而减少kernel的launch间隙时间。

cuda_graph基础

根据官方的源码示例，对cuda_graph的示例进行了补充，以下是完整代码示例：

#include <iostream>
#include <cuda_stream.h>

const int N = 5000;
const int NSTEP = 1000;
const int NKERNEL = 20;

const int blocks = 32;
const int threads = 512;

__global__ void shortKernel(float *out_d, float *in_d){
  int idx=blockIdx.x*blockDim.x+threadIdx.x;
  if(idx<N){
    out_d[idx]=1.5*in_d[idx];
  }
}

void run1(cudaStream_t &stream, float *out_d, float *in_d){
    // start CPU wallclock timer
    for(int istep=0; istep<NSTEP; istep++){
        for(int ikrnl=0; ikrnl<NKERNEL; ikrnl++){
            shortKernel<<<blocks, threads, 0, stream>>>(out_d, in_d);
            cudaStreamSynchronize(stream);
        }
    }
}

void run2(cudaStream_t &stream, float *out_d, float *in_d){
    // start wallclock timer
    for(int istep=0; istep<NSTEP; istep++){
        for(int ikrnl=0; ikrnl<NKERNEL; ikrnl++){
            shortKernel<<<blocks, threads, 0, stream>>>(out_d, in_d);
        }
        cudaStreamSynchronize(stream);
    }
    //end wallclock timer
}

void run3(cudaStream_t &stream, float *out_d, float *in_d){
    bool graphCreated=false;
    cudaGraph_t graph;
    cudaGraphExec_t instance;
    for(int istep=0; istep<NSTEP; istep++){
        if(!graphCreated){
            cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
            for(int ikrnl=0; ikrnl<NKERNEL; ikrnl++){
                shortKernel<<<blocks, threads, 0, stream>>>(out_d, in_d);
            }
            cudaStreamEndCapture(stream, &graph);
            cudaGraphInstantiate(&instance, graph, NULL, NULL, 0);
            graphCreated=true;
        }
    cudaGraphLaunch(instance, stream);
    cudaStreamSynchronize(stream);
    }
}

int main(int argc, char const *argv[])
{
    /* code */
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    float* in_h = new float[N];
    float* out_h = new float[N];
    int nBytes = N * sizeof(float);
    // float *in_h, *out_h;
    // in_h = (float*)malloc(nBytes);
    // out_h = (float*)malloc(nBytes);

    for (int i=0;i<N;i++){
        in_h[i] = i+1;
    }
    
    float *in_d, *out_d;
    cudaMalloc((void **)&in_d, nBytes);
    cudaMalloc((void **)&out_d, nBytes);

    cudaMemcpy(in_d, in_h, nBytes, cudaMemcpyHostToDevice);

    // run1(stream, out_d, in_d);
    // run2(stream, out_d, in_d);
    run3(stream, out_d, in_d);

    cudaMemcpy(out_h, out_d, nBytes, cudaMemcpyDeviceToHost);

    for (int i=0;i<N;i++){
        printf("%f\n", out_h[i]);
    }

    cudaFree(in_d);
    cudaFree(out_d);
    delete[] in_h;
    delete[] out_h;

    return 0;
}

这里run1对应的是每个kernel都要进行同步，两个kernel之间会有大量的launch时间，对应的profile就像下图一样，可以看到每个kernel之间的gap时间是大于kernel本身的计算时间的，这就带来了极大的开销。

)

这里run2对应的是每20次kernel进行一次同步，也即视为每20个kernel相互间是独立的，这样20个kernel间便可以不等上一个kernel结束便开始下一轮启动，中间就可以overlap掉一部分开销，不过在同步之间的gap依然很大，profile图示如下：

这里run2对应的是将每20个小kernel进行capture得到一个cuda_graph，然后在后续的每次运行时只需launch这个cuda_graph即可，对于这20个小kernel，便是连续的没有中间启动时间开销的，这样就可以加快运行速度，其profile图示如下所示：

controlnet-cuda_graph优化

理解了上面cuda_graph运行原理，那么便不难写出优化代码了，具体如下：

with open('trt/{}.plan'.format(model), 'rb') as f, trt.Runtime(trt_logger) as runtime:
            trt_engine = runtime.deserialize_cuda_engine(f.read())
            trt_ctx = trt_engine.create_execution_context()

            for index in range(trt_engine.num_io_tensors):
                name = trt_engine.get_binding_name(index)
                if 'vae' in model:
                    if name == 'z':
                        trt_ctx.set_tensor_address(name, self.tensors['x'].data_ptr())
                    if name == 'out':
                        trt_ctx.set_tensor_address(name, self.tensors['img_out'].data_ptr())
                else:
                    trt_ctx.set_tensor_address(name, self.tensors[name].data_ptr())

            trt_ctx.execute_async_v3(stream)
            
            if useGraph:
                cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)
                trt_ctx.execute_async_v3(stream)
                graph = cudart.cudaStreamEndCapture(stream)[1]
                graph_instance = cudart.cudaGraphInstantiate(graph, 0)[1]
            else:
                graph_instance = None

            self.engine_context_map[model] = trt_ctx

关键在于在每次运行前都capture对应的stream，形成graph_instance，而在具体使用时，只需要将对应的graph_instance启动起来即可：

 if useGraph:
    cudart.cudaGraphLaunch(self.control_fp16_graph_instance, self.stream1)
    cudart.cudaEventRecord(self.event1, self.stream1)
    cudart.cudaGraphLaunch(self.unet_input_fp16_graph_instance, self.stream)
    cudart.cudaStreamWaitEvent(self.stream, self.event1, cudart.cudaEventWaitDefault)
    cudart.cudaGraphLaunch(self.unet_output_fp16_graph_instance, self.stream)
else:
    self.engine_context_map['sd_control_fp16'].execute_async_v3(self.stream1)
    cudart.cudaEventRecord(self.event1, self.stream1)
    self.engine_context_map['sd_unet_input_fp16'].execute_async_v3(self.stream)
    cudart.cudaStreamWaitEvent(self.stream, self.event1, cudart.cudaEventWaitDefault)
    self.engine_context_map['sd_unet_output_fp16'].execute_async_v3(self.stream)

Multi-stream

cudaStream分为隐式流(默认流)和显式流。

对于隐式流：

所有的CUDA操作默认运行在隐式流里；
隐式流里的GPU 操作和CPU 操作两者是同步的；

对于显式流：

CPU计算和kernel计算并行；
CPU计算和数据传输并行；
数据传输和kernel计算并行
不同显式流的kernel计算并行

对于controlnet,可以将clip编码部分和vae的解码部分分割维两个不同的流，这样就可以利用输入输出的pipeline，上一个数据的解码便可以与下一个数据的编码同时进行。对于Unet和control部分，也可以通过分割为两个不同的stream进行处理。

self.stream = cudart.cudaStreamCreateWithPriority(cudart.cudaStreamNonBlocking, 0)[1]  
self.stream1 = cudart.cudaStreamCreateWithPriority(cudart.cudaStreamNonBlocking, 0)[1]
self.event = cudart.cudaEventCreateWithFlags(cudart.cudaEventDisableTiming)[1]
self.event1 = cudart.cudaEventCreateWithFlags(cudart.cudaEventDisableTiming)[1]

self.clip_instance = self.load_engine('sd_clip', self.stream1)
self.vae_instance = self.load_engine('sd_vae_fp16', self.stream)

self.control_fp16_graph_instance = self.load_engine('sd_control_fp16', self.stream1)
self.unet_input_fp16_graph_instance = self.load_engine('sd_unet_input_fp16', self.stream)
self.unet_output_fp16_graph_instance = self.load_engine('sd_unet_output_fp16', self.stream)

在数据同步上，通过cudaEventRecord标记同步点，通过cudaStreamWaitEvent同步时间，其实是同步不同stream之间的数据，保持数据一致性，下面的例子中clip模型要等待前面tokenizer操作完成，才能进行下操作，保持stream1开始前所有stream流中的数据都处理好了，重叠的是第二次的tokenizer和第一次的clip操作，这样整个运行是个流水线。最后通过cudaStreamSynchronize来保持所有同步。

cudart.cudaEventRecord(self.event, self.stream)
cudart.cudaStreamWaitEvent(self.stream1, self.event, cudart.cudaEventWaitDefault)
if useGraph:
    cudart.cudaGraphLaunch(self.clip_instance, self.stream1)
else:
    self.engine_context_map['sd_clip'].execute_async_v3(self.stream1)

cudart.cudaEventRecord(self.event1, self.stream1)