Make flashinfer kernels cuda graphs friendly about flashinfer HOT 7 OPEN

#include <cuda_runtime.h>
#include <iostream>


__global__ void subKernel(int *data) {
    printf("Data before sub kernel: %d\n", *data);
    (*data) -= 1;
}

__global__ void addKernel(int *data) {
    printf("Data before add kernel: %d\n", *data);
    (*data) += 1;
}

struct UserData {
    int data;
    bool op;
};

__global__ void launchChildKernelFromDevice(void *_userData) {
    UserData *userData = (UserData *)_userData;
    bool op = userData->op;

    if (op) {
        addKernel<<<1, 1>>>((int*)userData);
    } else {
        subKernel<<<1, 1>>>((int*)userData);
    }
}

int main() {
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    UserData *userData;
    cudaMallocHost(&userData, sizeof(UserData));

    userData->data = 10;
    userData->op = true;

    // run add kernel for sanity check

    cudaStreamSynchronize(stream);
    std::cout << "Data before kernel: " << userData->data << std::endl;
    launchChildKernelFromDevice<<<1, 1, 0, stream>>>(userData);
    cudaStreamSynchronize(stream);
    std::cout << "Data after kernel: " << userData->data << std::endl;

    cudaGraph_t graph;
    cudaGraphExec_t instance;

    // Begin graph capture
    cudaStreamSynchronize(stream);
    cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);

    // Use cuda host function to launch child kernel
    launchChildKernelFromDevice<<<1, 1, 0, stream>>>(userData);

    // End graph capture
    cudaStreamEndCapture(stream, &graph);
    cudaGraphInstantiate(&instance, graph, NULL, NULL, 0);
    
    cudaStreamSynchronize(stream);

    printf("Data after graph: %d\n", userData->data);

    // Run the graph
    cudaGraphLaunch(instance, stream);
    cudaStreamSynchronize(stream);

    printf("Data after graph replay: %d\n", userData->data);

    userData->op = false;
    cudaGraphLaunch(instance, stream);
    cudaStreamSynchronize(stream);

    printf("Data after graph replay with different op: %d\n", userData->data);

    cudaGraphExecDestroy(instance);
    cudaGraphDestroy(graph);
    cudaStreamDestroy(stream);
    cudaFree(userData);

    return 0;
}